# io = open("C:/Users/Andrew.Bannerman/Desktop/Julia/My_First_Neural_Network/Data/train-images-idx3-ubyte/train-images.idx3-ubyte", n=28)
#write(io)

#open("C:/Users/Andrew.Bannerman/Desktop/Julia/My_First_Neural_Network/Data/train-images-idx3-ubyte/train-images.idx3-ubyte") do file
#    for ln in eachline(file)
#        println("$(length(ln)), $(ln)")
#    end
#end

# Back propogation neural network
# 'Hello World' of neural networks
# Andrew Bannerman 4.21.2019

# Resources
# Cheat Sheet = https://ml-cheatsheet.readthedocs.io/en/latest/forwardpropagation.html
# Matrix multiplication Rules - https://www.mathsisfun.com/algebra/matrix-multiplying.html
# How to get MNIST data as a CSV https://pjreddie.com/projects/mnist-in-csv/
# Kaggle - https://www.kaggle.com/ngbolin/mnist-dataset-digit-recognizer

# Back propogation tutorial: https://dev.to/shamdasani/build-a-flexible-neural-network-with-backpropagation-in-python

# Needed packages
using DataFrames
using CSV
using Distributions
using Plots
using Random

#=
Random.seed!(123) # Setting the seed
d = Normal()
x = rand(d, 100)
x = rand(Normal(100, 2), 100)
vars = var(x)
sqrt(vars)
histogram(x)
=#

# Train and test set data preparation

# Load MNIST Data set from .csv format
train_set = CSV.read("C:/Users/Andrew.Bannerman/Desktop/Julia/My_First_Neural_Network/Data/mnist_train.csv",header=false,types=fill(Float64,785))
test_set = CSV.read("C:/Users/Andrew.Bannerman/Desktop/Julia/My_First_Neural_Network/Data/mnist_test.csv",header=false,types=fill(Float64,785))

# Plot Example image
example = Float64.(test_set[2,2:size(test_set,2)])
example = rotl90(reshape(example,28,28))
heatmap(example,c = :greys)
print(example)

# Extract labels
train_label = train_set[:,1]
test_label = test_set[:,1]

# Extract Data and exclude the label in the first column
train_set = train_set[:,2:size(train_set,2)]
test_set = test_set[:,2:size(test_set,2)]

# Pre process - Scale the inputs to a range of 0.01 to 1.00
m_train = Array{Float64}(zeros(nrow(train_set),size(train_set,2))) # empty matrix
@inbounds for i in 1:nrow(train_set)
    @inbounds for j in 1:size(train_set,2)
    m_train[i,j] = (train_set[i,j] / 255.0 * 0.99) + 0.01
    end
end

m_test = Array{Float64}(zeros(nrow(test_set),size(test_set,2))) # empty matrix
@inbounds for i in 1:nrow(test_set)
    @inbounds for j in 1:size(test_set,2)
    m_test[i,j] = (test_set[i,j] / 255.0 * 0.99) + 0.01
    end
end

# Convert matrix to DataFrame and append the labels back to the train and test set
scaled_orig_train_inputs = DataFrame(hcat(train_label,m_train))
scaled_test_inputs = DataFrame(hcat(test_label,m_test))

# Plot the new scale example
example = Float64.(scaled_orig_train_inputs[1,2:size(scaled_orig_train_inputs,2)])
example = rotl90(reshape(example,28,28))
heatmap(example, c = :greys)

#---------
# Data Augmentation
#---------

# Rotate images
training_set_aug_minus = zeros(size(scaled_orig_train_inputs,1),size(scaled_orig_train_inputs,2))
training_set_aug_plus = zeros(size(scaled_orig_train_inputs,1),size(scaled_orig_train_inputs,2))
degree_rot_minus = -.3
degree_rot_plus = .3
i=1
for i in 1:size(scaled_orig_train_inputs,1)
        label = scaled_orig_train_inputs[i,1]
        example = Float64.(scaled_orig_train_inputs[i,2:size(scaled_orig_train_inputs,2)])
        example = reshape(example,28,28)
        #rot_minus = LinearMap(RotMatrix(degree_rot_minus))
        #rot_plus = LinearMap(RotMatrix(degree_rot_plus))
        #imgw_minus = warp(example_minus, rot_minus, axes(example_plus)) #  axes(example_minus) = crops
        #imgw_plus = warp(example_plus, rot_plus, axes(example_plus)) # axes() = crops
        imgw_minus = imrotate(example, degree_rot_minus, axes(example))
        imgw_plus = imrotate(example, degree_rot_plus, axes(example))
        imgw_minus = reshape(imgw_minus,1,784)
        imgw_plus = reshape(imgw_plus,1,784)
            for k in 1:size(imgw_minus,2)
                if isnan(imgw_minus[k]) # Convert NaN to 0.0
                    imgw_minus[k] = 0.01
                end
                if isnan(imgw_plus[k]) # Convert NaN to 0.0
                    imgw_plus[k] = 0.01
                end
            end
            #imgw_plus .= (imgw_plus ./ 255.0 .* 0.99) .+ 0.01 # Scale the outputs
            #imgw_minus .= (imgw_minus ./ 255.0 .* 0.99) .+ 0.01 # Scale the outputs
            training_set_aug_minus[i,1] = label
            training_set_aug_minus[i,2:size(scaled_orig_train_inputs,2)] = imgw_minus
            training_set_aug_plus[i,1] = label
            training_set_aug_plus[i,2:size(scaled_orig_train_inputs,2)] = imgw_plus
        print("This is iteration i",i,"\n")
end

    maximum(imgw_plus)

# Combine augmented data with the original
# note may want to add shuffling?
scaled_train_inputs = DataFrame(vcat(scaled_orig_train_inputs,training_set_aug_minus,training_set_aug_plus))
#scaled_train_inputs = DataFrame(training_set_aug_plus)

# Shuffle the augmented + original data
scaled_train_inputs = scaled_train_inputs[shuffle(1:end), :]

#---------
# Activation Functions
#---------

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

function sigmoid_derivative(x)
    s = 1 / (1 + exp(-x))
    return s * (1 - s)
end

function ReLU(x)
    #return x*(x>0)
    return maximum([0.0,x])
end

function ReLU_derivative(x)
    if x <= 0
    end
    return 0
end

#---------
# Layers
# https://medium.freecodecamp.org/building-a-3-layer-neural-network-from-scratch-99239c4af5d3
#---------

function dropout(d_in::Array{Float64,2}; p::Float64=.8)
    out = (d_in .* rand(Binomial(1,p), size(d_in)[1])) ./ p
return out
end

activation = "sigmoid"
drop_out = false
function forward_prop(W1::Array{Float64,2}, b1::Array{Float64,2}, W2::Array{Float64,2}, b2::Array{Float64,2}, X::Array{Float64,2}, Y::Array{Float64,2}; drop_out::Bool=true, p::Float64=.5, activation::String="sigmoid")
    Z1 = W1 * X + b1
    if drop_out == false
        if activation == "sigmoid"
    A1 = sigmoid.(Z1)
elseif activation == "relu"
    A1 = ReLU.(Z1)
    end
    elseif drop_out == true
        if activation == "sigmoid"
    A1 = sigmoid.(Z1)
    A1 = dropout(A1,p=.8)
    elseif activation == "relu"
    A1 = ReLU.(Z1)
    A1 = dropout(A1,p=.8) # Apply drop out by hidden output * 0,1's / p (1/p - scaled so we dont need to apply at test time)(inverted dropout)
        end
    end
    Z2 = (W2 * A1) + b2
    if activation == "sigmoid"
    A2 = sigmoid.(Z2)
elseif activation == "relu"
    A2 = ReLU.(Z2)
    end
    return Z1, A1, Z2, A2
end

#----
# back prop math: http://deeplizard.com/learn/video/G5b4jRBKNxw
# https://github.com/mbadry1/DeepLearning.ai-Summary/tree/master/1-%20Neural%20Networks%20and%20Deep%20Learning#gradient-descent-on-m-examples
#----
function backward_prop(W1::Array{Float64,2}, b1::Array{Float64,2}, W2::Array{Float64,2}, b2::Array{Float64,2}, X::Array{Float64,2}, Y::Array{Float64,2}, Z1::Array{Float64,2}, A1::Array{Float64,2}, Z2::Array{Float64,2}, A2::Array{Float64,2}; learning_rate::Float64=.01,activation::String="relu")
    dZ2 = A2 - Y # output errors
    dW2 = (dZ2 * A1')
    db2 =  sum(dZ2, dims = 2)
    if activation == "sigmoid"
        dZ1 = (W2' * dZ2) .* sigmoid_derivative.(A1) #(1 .- A1.^2)
    elseif activation == "relu"
        dZ1 = (W2' * dZ2) .* ReLU_derivative.(A1) #(1 .- A1.^2)
    end
    dW1 = (dZ1 * X')
    db1 =  sum(dZ1,dims = 2)
   # Update params
    W1 = W1 .- (learning_rate .* dW1)
     b1 = b1 .- (learning_rate .* db1)
     W2 = W2 .- (learning_rate * dW2)
     b2 = b2 .- (learning_rate .* db2)
    # print("check w1 is updating",W2[1])

    return  dW1, db1, dW2, db2, W1, b1, W2, b2
end

#---------
# Cost Function
#---------

function cost_function(A2::Array{Float64,2}, Y::Array{Float64,2}; activation::String="sigmoid")
    if activation == "sigmoid"
    m = size(Y)[1]
    #cost = (- 1 / m) * np.sum(Y * np.log(A) + (1 - Y) * (np.log(1 - A)))  # compute cost
    cost = (-1/m) * sum((Y' * log.(A2)) .+ ((1 .- Y')* log.(1 .- A2)))
        #cost = (-1/m) * sum((Y * log.(A2')) .+ ((1 .- Y)* log.(1 .- A2')))
elseif activation == "relu"
    return Inf
end
return cost
end

#---------
# Initialize network
#---------

# Define neural network
# Number of input, hidden and output nodes
input_nodes = size(train_set,2) # equal to the size of the 28 * 28 image (784 elements)
hidden_nodes = 256
output_nodes = 10 # Equal to the number of labels in this case numbers 0 to 9 (10 elements total)

# Create instance of neural network
inodes = input_nodes
hnodes = hidden_nodes
onodes = output_nodes
# Weight initialization
# w = np.random.randn(n) * sqrt(2.0/n) for relu
# w = np.random.randn(n) / sqrt(n) # other
#W1 = randn(hnodes,inodes) * inodes^(-0.5)
#W2 = randn(onodes,hnodes) * hnodes^(-0.5)

W1 = randn(hnodes,inodes) * sqrt(2.0/inodes)
W2 = randn(onodes,hnodes) * sqrt(2.0/hnodes)
# Initializing the biases as 0's
b1 = zeros(hnodes,1)
b2 = zeros(onodes,1)

epochs = 5
# Check that the neural network is converging
# Initialize training progress
learning_curve = fill(0,nrow(scaled_train_inputs))
training_accuracy = fill(0.0,nrow(scaled_train_inputs))
i = 1
j = 1
@inbounds for i in 1:epochs
    # go through all records in the training data set
    @inbounds for j in 1:nrow(scaled_train_inputs)
        # Forward propogation
        label = Int64.(scaled_train_inputs[j,1])
        X = reshape(Float64.(scaled_train_inputs[j,2:size(scaled_train_inputs,2)]),784,1)
        # create the target output values (all 0.01, except the desired label which is 0.99)
        Y = zeros(10,1)
        Y .= 0.01 .+ Y
        # Set the target .99 to the correct index position within the target array
        Y[label+1] = 0.99 # We +1 because the number convention starts at 0 to 9. Thus array position 6 is actually target number 5.
global Z1, A1, Z2, A2 = forward_prop(W1, b1, W2, b2, X, Y; drop_out=false, p=.5,activation="relu")
cost_rate = cost_function(A2,Y;activation="relu")
global dW1, db1, dW2, db2, W1, b1, W2, b2 = backward_prop(W1, b1, W2, b2, X, Y, Z1, A1, Z2, A2; learning_rate=.01,activation="relu")
#print("check training weight being update", dW1[1],"\n")
network_train_output = argmax(A2)[1]-1 # adjust index position to account from 0 start
if (label == network_train_output)
    # network's answer matches correct answer, add 1 to learning curve
    learning_curve[j] = 1
else
    # network's answer doesn't match correct answer, add 0 to learning curve
    learning_curve[j] = 0
end
if (j == 1)
    global learning_curve = fill(0,nrow(scaled_train_inputs))
end
print("Cost after iteration ",j," ",cost_rate,"\n")
print("Training Accuracy ",(sum(learning_curve) / j),"\n")
print("This is epoch ",i,"\nThis is iteration ",j,"\n")

    end
end

sum(learning_curve) / 60000

print(learning_curve)

# Plot the learning curve
plot(cumsum(learning_curve),title="Learning Curve")

#---------
# Test Function
#---------

# Function for applying the network to unseen or test data not seen (trained) by the network before
function test_nn(W1::Array{Float64,2}, W2::Array{Float64,2},b1::Array{Float64,2}, b2::Array{Float64,2},X::Array{Float64,2})
    Z1 = W1 * X + b1
        if activation == "sigmoid"
    A1 = sigmoid.(Z1)
elseif activation == "relu"
    A1 = ReLU.(Z1)
    end
    Z2 = (W2 * A1) + b2
    if activation == "sigmoid"
    A2 = sigmoid.(Z2)
elseif activation == "relu"
    A2 = ReLU.(Z2)
    end
    return A2
end

#---------
# Test Network on Individual image
#---------

# Call the network on an indivdual entry
example_call = Float64.(scaled_test_inputs[500,2:size(scaled_test_inputs,2)])
example_call = reshape(example_call,784,1)
example_call_image = rotl90(reshape(example_call,28,28))
# Plot Number
heatmap(example_call_image,c = :greys)
# Grab the known correct label
correct_label = scaled_test_inputs[500,1]
# Apply the trained network to the input example
outputs = test_nn(W1, W2, b1,b2,example_call)
# Check what the label thought the number was
network_label = argmax(outputs)[1]-1
if network_label ==  correct_label
    print("The neural network correctly identified the hand written number - ", network_label)
else
    print("The neural network incorrectly identified the hand written number")
end

print(example_call_image)

#---------
# Test Neural Network on Test Data
#---------

# Initialize the scorecard output the same size as the number of test set labels
scorecard = fill(0,nrow(scaled_test_inputs))
test_accuracy = fill(0.0,nrow(scaled_test_inputs))

# Input the test set data over the trained neural network to see how well it does
i=1
@inbounds for i in 1:nrow(scaled_test_inputs)
    let wih = wih, who = who
    correct_label = Int64.(scaled_test_inputs[i,1])
    inputs = reshape(Float64.(scaled_test_inputs[i,2:size(scaled_test_inputs,1)]),784,1) # rotr90() for changing dim to [784,1] from [1,784] to meet matrix multiplication rules
    # query the network
    outputs = test_nn(W1, W2, b1,b2,inputs) # Test the network over the test set data (wih and who already set during training - no weight updating happens this time round)
    # Find element position of what the network thinks the output is ie if output is a .96 in element position 6 then the network thinks the label was 5.
    # This is because we have hand written numbers 0 to 9, total 10 outputs. Starting from index position 1 = label 0 hence index position 6 then = label 5
    label = argmax(outputs)[1]-1 # adjust index position to account from 0 start
    # append correct or incorrect to list
    if (label == correct_label)
        # network's answer matches correct answer, add 1 to scorecard
        scorecard[i] = 1
    else
        # network's answer doesn't match correct answer, add 0 to scorecard
        scorecard[i] = 0
        end
        print("Accuracy ",sum(scorecard) / i,"\n")
        test_accuracy[i] = sum(scorecard) / i
    end
end

# Percentage Correct
percentage = sum(scorecard) / size(scorecard,1)

# Change log
# 5.12.2019 - Changed the back propogation derivatives to allow different activation functions
# 5.12.2019 - Added bias
# 4.23.2019 - Added first implemntation of drop out
# 4.22.2019 - added a % print function to check train / test accuracy


using Random, Distributions, Plots

function cum_miss(; nobs, nvar, p)
    b = Bernoulli(p)
    sample = rand(b, nobs*nvar)
    sample = reshape(sample, (nobs, nvar))
    patterns = [sample[i, :] for i in 1:size(sample, 1)]

    #binomial = sum.(patterns)
    #histogram(binomial)

    remember = Vector{BitArray{1}}()
    rows = [Vector{Int64}(undef, 0) for i = 1:size(patterns, 1)]

    for i = 1:size(patterns, 1)
        unknown = true
        for j = 1:size(remember, 1)
            if patterns[i] == remember[j]
                push!(rows[j], i)
                unknown = false
            end
        end
        if unknown
            push!(remember, patterns[i])
            push!(rows[size(remember, 1)], i)
        end
    end

    rows = rows[1:length(remember)]
    n_pattern = size(rows, 1)

    # sort by number of missings
    sort_n_miss = sortperm(sum.(remember))
    remember = remember[sort_n_miss]
    remember_cart = findall.(!, remember)
    rows = rows[sort_n_miss]

    pattern_n_obs = size.(rows, 1)
    pattern_nvar_obs = length.(remember_cart)

    obs_pat = pattern_n_obs.*pattern_nvar_obs
    sort_n_miss = sortperm(obs_pat, rev = true)
    obs_pat = obs_pat[sort_n_miss]
    pattern_n_obs = pattern_n_obs[sort_n_miss]

    #pattern_n_obs_cum = pattern_n_obs_cum[sort_n_miss]
    pattern_nvar_obs = pattern_nvar_obs[sort_n_miss]
    #complexity = complexity[sort_n_miss]

    pattern_n_obs_cum = cumsum(pattern_n_obs)

    complexity = pattern_nvar_obs.^3
    obs_pat_cum = cumsum(obs_pat)

    p1 = plot(1:n_pattern, obs_pat_cum, title = "observed datapoints cum.", ylims = (0, Inf))
    p2 = plot(1:n_pattern, pattern_n_obs, title = "N observed")
    p3 = plot(1:n_pattern, pattern_n_obs_cum, title = "N observed cum.", ylims = (0, Inf))
    p4 = plot(1:n_pattern, pattern_nvar_obs, title = "N variables")
    p5 = plot(1:n_pattern, complexity, title = "N variables ^3")

    plot(p1, p2, p3, p4, layout = (2,2))
end

cum_miss(;nobs= 1000, nvar= 10, p= 0.1)

cum_miss(;nobs= 1000, nvar= 100, p= 0.1)

cum_miss(;nobs= 5000, nvar= 100, p= 0.1)

cum_miss(;nobs= 1000, nvar= 100, p= 0.01)

cum_miss(;nobs= 10000, nvar= 100, p= 0.01)

cum_miss(;nobs= 1000, nvar= 10, p= 0.9)

a = [Base.binomial(100, k) for k = 0:5]

function sim_dropout(nwaves, dist)
    pattern = falses(nwaves)
    sample = rand(dist, nwaves)
    t_dropout = findfirst(sample)
    !isnothing(t_dropout) ?
        pattern[t_dropout:nwaves] .= true :
        nothing
    pattern[1] = false
    return pattern
end

function get_miss(nvar, dist, miss_wave)
    if miss_wave
        pattern = trues(nvar)
    else
        pattern = rand(dist, nvar)
    end
    return pattern
end

function pattern_long(nvar, dist, miss_waves)
    pattern =
        [get_miss(nvar, dist, miss_wave) for miss_wave in miss_waves]
    pattern = vcat(pattern...)
    return(pattern)
end

function cum_miss_dropout(; nobs, nvar, nwaves, p, p_dropout)
    dist1 = Bernoulli(p_dropout)
    dist2 = Bernoulli(p)
    miss_waves = [sim_dropout(nwaves, dist1) for i in 1:nobs]

    patterns = [pattern_long(nvar, dist2, miss_waves[i]) for i in 1:nobs]

    #return patterns
    #binomial = sum.(patterns)
    #histogram(binomial)

    remember = Vector{BitArray{1}}()
    rows = [Vector{Int64}(undef, 0) for i = 1:size(patterns, 1)]

    for i = 1:size(patterns, 1)
        unknown = true
        for j = 1:size(remember, 1)
            if patterns[i] == remember[j]
                push!(rows[j], i)
                unknown = false
            end
        end
        if unknown
            push!(remember, patterns[i])
            push!(rows[size(remember, 1)], i)
        end
    end

    rows = rows[1:length(remember)]
    n_pattern = size(rows, 1)

    # sort by number of missings
    sort_n_miss = sortperm(sum.(remember))
    remember = remember[sort_n_miss]
    remember_cart = findall.(!, remember)
    rows = rows[sort_n_miss]

    pattern_n_obs = size.(rows, 1)
    pattern_nvar_obs = length.(remember_cart)

    obs_pat = pattern_n_obs.*pattern_nvar_obs
    sort_n_miss = sortperm(obs_pat, rev = true)
    obs_pat = obs_pat[sort_n_miss]
    pattern_n_obs = pattern_n_obs[sort_n_miss]

    #pattern_n_obs_cum = pattern_n_obs_cum[sort_n_miss]
    pattern_nvar_obs = pattern_nvar_obs[sort_n_miss]
    #complexity = complexity[sort_n_miss]

    pattern_n_obs_cum = cumsum(pattern_n_obs)

    complexity = pattern_nvar_obs.^3
    obs_pat_cum = cumsum(obs_pat)

    p1 = plot(1:n_pattern, obs_pat_cum, title = "observed datapoints cum.", ylims = (0, Inf))
    p2 = plot(1:n_pattern, pattern_n_obs, title = "N observed")
    p3 = plot(1:n_pattern, pattern_n_obs_cum, title = "N observed cum.", ylims = (0, Inf))
    p4 = plot(1:n_pattern, pattern_nvar_obs, title = "N variables")
    p5 = plot(1:n_pattern, complexity, title = "N variables ^3")

    plot(p1, p2, p3, p4, layout = (2,2))
end

cum_miss_dropout(nobs = 1000, nvar = 10, nwaves = 10, p = 0.05, p_dropout = 0.1)

dist = Binomial(10, 0.1)


# Ground truth was generated using the bw.SJ() R function, with nb=10000 and tol=1e-3 unless further specified.
# Some differences are expected because of how the double sums are approximated.

@testset "fewpoints" begin
	@test isnan(bwsj([0.0]))
	@test isnan(bwsj([1.0]))
	@test isnan(bwsj(ones(100)))
	@test bwsj([1.0,2.0])≈0.1163591 rtol=0.1
end

@testset "range" begin
	# R: nb=10000000, tol=1e-6
	@test bwsj(range(0,stop=1,length= 96); rtol=1e-1)≈0.1109613 rtol=1e-1
	@test bwsj(range(0,stop=1,length= 97); rtol=1e-1)≈0.1105281 rtol=1e-1
	@test bwsj(range(0,stop=1,length= 98); rtol=1e-1)≈0.1101011 rtol=1e-1
	@test bwsj(range(0,stop=1,length= 99); rtol=1e-1)≈0.1096801 rtol=1e-1
	@test bwsj(range(0,stop=1,length=100); rtol=1e-1)≈0.1092651 rtol=1e-1
	@test bwsj(range(0,stop=1,length=101); rtol=1e-1)≈0.1088558 rtol=1e-1
	@test bwsj(range(0,stop=1,length=102); rtol=1e-1)≈0.1084521 rtol=1e-1
	@test bwsj(range(0,stop=1,length=103); rtol=1e-1)≈0.1080540 rtol=1e-1
	@test bwsj(range(0,stop=1,length=104); rtol=1e-1)≈0.1076612 rtol=1e-1
	@test bwsj(range(0,stop=1,length=105); rtol=1e-1)≈0.1072737 rtol=1e-1

	@test bwsj(range(0,stop=1,length= 96); rtol=1e-3)≈0.1109613 rtol=1e-3
	@test bwsj(range(0,stop=1,length= 97); rtol=1e-3)≈0.1105281 rtol=1e-3
	@test bwsj(range(0,stop=1,length= 98); rtol=1e-3)≈0.1101011 rtol=1e-3
	@test bwsj(range(0,stop=1,length= 99); rtol=1e-3)≈0.1096801 rtol=1e-3
	@test bwsj(range(0,stop=1,length=100); rtol=1e-3)≈0.1092651 rtol=1e-3
	@test bwsj(range(0,stop=1,length=101); rtol=1e-3)≈0.1088558 rtol=1e-3
	@test bwsj(range(0,stop=1,length=102); rtol=1e-3)≈0.1084521 rtol=1e-3
	@test bwsj(range(0,stop=1,length=103); rtol=1e-3)≈0.1080540 rtol=1e-3
	@test bwsj(range(0,stop=1,length=104); rtol=1e-3)≈0.1076612 rtol=1e-3
	@test bwsj(range(0,stop=1,length=105); rtol=1e-3)≈0.1072737 rtol=1e-3
end

@testset "reference" begin
	@test bwsj_reference(range(0,stop=1,length=100); rtol=1e-1)≈0.1092651 rtol=1e-1
end

@testset "scale" begin
	@test bwsj(range(0,stop=1e3,  length=101))≈0.1088518*1e3   rtol=0.1
	@test bwsj(range(0,stop=1e-3, length=101))≈0.1088518*1e-3  rtol=0.1
	@test bwsj(range(0,stop=1e24, length=101))≈0.1088518*1e24  rtol=0.1
	@test bwsj(range(0,stop=1e-24,length=101))≈0.1088518*1e-24 rtol=0.1
end

@testset "randn" begin
	rng = StableRNG(1234)
	# @info "A"
	X = randn(rng, 20)
	@test bwsj(X)≈0.6189722 rtol=0.1

	# @info "B"
	X = randn(rng, 100)
	@test bwsj(X)≈0.4314236 rtol=0.1

	# @info "C"
	X = randn(rng, 1000)
	@test bwsj(X)≈0.2823561 rtol=0.1
end

@testset "quantile" begin
	N = 100
	X = quantile.(Normal(), range(1/2N,stop=1-1/2N,length=N))

	@test bwsj(sign.(X).*abs.(X).^(2.0.^-4))≈0.05301544  rtol=0.1
	@test bwsj(sign.(X).*abs.(X).^(2.0.^-3))≈0.07877113  rtol=0.1
	@test bwsj(sign.(X).*abs.(X).^(2.0.^-2))≈0.1279672   rtol=0.1
	@test bwsj(sign.(X).*abs.(X).^(2.0.^-1))≈0.2092015   rtol=0.1
	@test bwsj(sign.(X).*abs.(X).^(2.0.^0 ))≈0.4747508   rtol=0.1
	@test bwsj(sign.(X).*abs.(X).^(2.0.^1 ))≈0.1394881   rtol=0.1
	@test bwsj(sign.(X).*abs.(X).^(2.0.^2 ))≈0.03132362  rtol=0.1
	@test bwsj(sign.(X).*abs.(X).^(2.0.^3 ))≈0.004705298 rtol=0.1 # NB: run with nb=10000000 to get accurate answer

	@test bwsj(sign.(X).*abs.(X).^(2.0.^6))≈1.1579829602522003e-12 rtol=0.1 # NB: ground truth generated with bwsj_reference(; rtol=1e-24)
end

@testset "outliers" begin
	# ground truth generated with bwsj_reference(; rtol=1e-24)
	rng = StableRNG(2017)
	X = rand(rng, 100)

	Y1 = 1e9 .+ randn(rng,1)
	@test bwsj(vcat(X,Y1))≈0.11706501880745836 rtol=0.1

	Y2 = 1e9 .+ randn(rng,10)
	@test bwsj(vcat(X,Y2))≈0.1197213803664104 rtol=0.1

	Y3 = 1e9.*sign.(randn(rng,10)) .+ randn(rng,10)
	@test bwsj(vcat(X,Y3))≈0.12247926574875453 rtol=0.1
end

@testset "repeats" begin
    X = vcat(zeros(10),ones(10),2*ones(10),4*ones(10))
    @test bwsj(X;leafsize=10)≈0.1246847 rtol=1
    @test bwsj(X;leafsize=5 )≈0.1246847 rtol=1
end

@testset "ϕbounds" begin
	for stepsize in (0.1, 0.5, 1.0, 2.5, 5.0)
		a = 0:stepsize:5
		b = a.+stepsize
		x = range(0,stop=stepsize,length=11)[2:end-1]'

		B = ϕ4bounds.(a,b,a.+x)
		lb,ub = first.(B),last.(B)
		@test all(lb.<=ϕ4.(a.+x).<=ub)

		B = ϕ6bounds.(a,b,a.+x)
		lb,ub = first.(B),last.(B)
		@test all(lb.<=ϕ6.(a.+x).<=ub)
	end
end


@testset "1035.uncrossed-lines.jl" begin
    @test max_uncrossed_lines([1, 4, 2], [1, 2, 4]) == 2
    @test max_uncrossed_lines([2, 5, 1, 2, 5], [10, 5, 2, 1, 5, 2]) == 3
    @test max_uncrossed_lines([1, 3, 7, 1, 7, 5], [1, 9, 2, 5, 1]) == 2
end


######### Jags line program example  ###########

using Mamba, Jags

ProjDir = joinpath(dirname(@__FILE__))
cd(ProjDir) do

  line = "
  model {
    for (i in 1:n) {
          mu[i] <- alpha + beta*(x[i] - x.bar);
          y[i]   ~ dnorm(mu[i],tau);
    }
    x.bar   <- mean(x[]);
    alpha    ~ dnorm(0.0,1.0E-4);
    beta     ~ dnorm(0.0,1.0E-4);
    tau      ~ dgamma(1.0E-3,1.0E-3);
    sigma   <- 1.0/sqrt(tau);
  }
  "

  monitors = Dict(
    "alpha" => true,
    "beta" => true,
    "tau" => true,
    "sigma" => true
    )

  jagsmodel = Jagsmodel(
    name="line1",
    model=line,
    monitor=monitors,
    deviance=true, dic=true, popt=true,
    pdir=ProjDir
    );

  println("\nJagsmodel that will be used:")
  jagsmodel |> display

  data = Dict(
    "x" => [1, 2, 3, 4, 5],
    "y" => [1, 3, 3, 3, 5],
    "n" => 5
  )

  inits = [
    Dict("alpha" => 0,"beta" => 0,"tau" => 1),
    Dict("alpha" => 1,"beta" => 2,"tau" => 1),
    Dict("alpha" => 3,"beta" => 3,"tau" => 2),
    Dict("alpha" => 5,"beta" => 2,"tau" => 5)
  ]

  println("Input observed data dictionary:")
  data |> display
  println("\nInput initial values dictionary:")
  inits |> display
  println()

  sim = jags(jagsmodel, data, inits, ProjDir)
  describe(sim)
  println()

  ## Brooks, Gelman and Rubin Convergence Diagnostic
  try
    gelmandiag(sim1, mpsrf=true, transform=true) |> display
  catch e
    #println(e)
    gelmandiag(sim, mpsrf=false, transform=true) |> display
  end

  ## Geweke Convergence Diagnostic
  gewekediag(sim) |> display

  ## Highest Posterior Density Intervals
  hpd(sim) |> display

  ## Cross-Correlations
  cor(sim) |> display

  ## Lag-Autocorrelations
  autocor(sim) |> display

  ## Deviance Information Criterion
  #dic(sim) |> display

  ## Plotting
  p = plot(sim, [:trace, :mean, :density, :autocor], legend=true);
  draw(p, nrow=4, ncol=4, filename="$(jagsmodel.name)-summaryplot", fmt=:svg)
  draw(p, nrow=4, ncol=4, filename="$(jagsmodel.name)-summaryplot", fmt=:pdf)

  # Below will only work on OSX, please adjust for your environment.
  # JULIA_SVG_BROWSER is set from environment variable JULIA_SVG_BROWSER
  @static is_apple() ? if isdefined(Main, :JULIA_SVG_BROWSER) && length(JULIA_SVG_BROWSER) > 0
          for i in 1:4
            isfile("$(jagsmodel.name)-summaryplot-$(i).svg") &&
              run(`open -a $(JULIA_SVG_BROWSER) "$(jagsmodel.name)-summaryplot-$(i).svg"`)
          end
        end : println()

  # Below examples of using other ways to display the simulation results
  (index, chains) = Jags.read_jagsfiles(jagsmodel)

  println()
  #chains[1]["samples"] |> display
  #println()

  if jagsmodel.dic
    (idx0, chain0) = Jags.read_pDfile(jagsmodel)
    #idx0 |> display
    println()
    #chain0[1]["samples"] |> display
  end

  if jagsmodel.dic || jagsmodel.popt
    pDmeanAndpopt = Jags.read_table_file(jagsmodel, data["n"]);
    pDmeanAndpopt["pD.mean"] |> display
    println()
    pDmeanAndpopt["popt"] |> display
  end

end #cd


include("DM1_1.jl")
include("DM2_1.jl")

function ReactiveGRASP(matrix,cost, n, m, ite, coupe,temps)
#    println("^^^^^^^^ Debut R-GRASP ^^^^^^^^")

    #initialisation
    (p,nb_iteration,z_cumul,zBest,zWorst) = intialiser(matrix,cost, n, m, ite, coupe)
    evol_p=Array{Float64}[]
#    println("valeurs initiales: ")
#    println("p: ", p)
#    println("nb_iteration: ", nb_iteration)
#    println("z_cumul: ", z_cumul)
#    println("zBest: ", zBest , "  zWorst: ", zWorst)

    t=time()
    while (time()-t <= temps)
        cpt = 1
        push!(evol_p,p)
        while (cpt <= ite)
            cpt = cpt+1

            prob=rand(Float64)

            alpha_choisit = choix_alpha(p,prob)
#            println("alpha_choisit: ", alpha_choisit, "  p: ", p)
            (SOL,z) = GRASP(cost, matrix, n, m, p[alpha_choisit])

            nb_iteration[alpha_choisit] += 1
            z_cumul[alpha_choisit] += z

            if z > zBest
                zBest = z
            end

            if z < zWorst
                zWorst = z
            end

#            println(" ============= \n ite: ",cpt)
        end
        println("zBest: ", zBest, "   zWorst: ", zWorst, "    nb_iteration : ", nb_iteration , "    p: ", p)
        recalcul_p!(p,z_cumul,zBest,zWorst,nb_iteration,evol_p)

    end
    println(evol_p)
end

#=======================================#

function intialiser(matrix,cost, n, m, ite, coupe)

    p = Float64[]
    for i in 1:coupe
        push!(p,i/coupe)
    end

    nb_iteration = zeros(Int64, length(p))

    z_cumul = zeros(Int64, length(p))

    (SOL,z) =Glouton(cost, matrix, n, m)

    zBest = z

    zWorst = z

    return(p,nb_iteration,z_cumul,zBest,zWorst)
end

#=======================================#

function recalcul_p!(p,z_cumul,zBest,zWorst,nb_iteration,evol_p)
    q=zeros(Float64,length(p))

    somme_q=0
#    println("q: ", q)
#    println("nb_iteration: ", nb_iteration)

    for i in 1:length(p)
#        println("q: ", q, "   i:", i)
#        if nb_iteration[i]!=0
#            println("z_cumul[i]: ", z_cumul[i],"   nb_iteration[i]: ", nb_iteration[i] )
            moyenne=(z_cumul[i]/nb_iteration[i])
#            println("moyenne: ", moyenne)
#        else
#            moyenne=0
#        end
#        if moyenne>0
            q[i]=(moyenne-zWorst)/(zBest-zWorst)
#            println("q[i]: ", q[i])

#        else
#            q[i]= zWorst/(zBest-zWorst)
#        end
        somme_q+=q[i]

    end
#    println("q: ", q, "      somme_q", somme_q)

    for i in 1:length(q)
        p[i]=q[i]/somme_q
    end
#    println("p: ", p)

end

#=======================================#

function choix_alpha(p,prob)
    for i in 1:length(p)
        if prob < p[i]
            return(i)
        elseif i==length(p)
            return(i)
        end
    end
end


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

"""
    Remainder([total])

The transform that takes a table with columns `x₁, x₂, …, xₙ`
and returns a new table with an additional column containing the
remainder value `xₙ₊₁ = total .- (x₁ + x₂ + ⋯ + xₙ)` If the `total`
value is not specified, then default to the maximum sum across rows.
"""
struct Remainder <: Transform
  total::Union{Float64,Nothing}
end

Remainder() = Remainder(nothing)

isrevertible(::Type{Remainder}) = true

assertions(::Type{Remainder}) = [TT.assert_continuous]

function _cache(transform::Remainder, table)
  # design matrix
  X = Tables.matrix(table)

  # find total across rows
  if !isnothing(transform.total)
    transform.total
  else
    maximum(sum(X, dims=2))
  end
end

function _apply(transform::Remainder, table, cache)
  # basic checks
  for assertion in assertions(transform)
    assertion(table)
  end

  # design matrix
  X = Tables.matrix(table)

  # retrieve the total
  total = cache

  # make sure that the total is valid
  @assert all(x -> x ≤ total, sum(X, dims=2)) "the sum for each row must be less than total"

  # original column names
  names = Tables.columnnames(table)

  # create a column with the remainder
  S = sum(X, dims=2)
  Z = [X (total .- S)]

  # create new column name
  rname = :remainder
  while rname ∈ names
    rname = Symbol(rname,:_)
  end
  names = (names..., rname)

  # table with new column
  𝒯 = (; zip(names, eachcol(Z))...)
  newtable = 𝒯 |> Tables.materializer(table)

  newtable, total
end

apply(transform::Remainder, table) =
  _apply(transform, table, _cache(transform, table))

function revert(::Remainder, newtable, cache)
  names = Tables.columnnames(newtable)
  TT.Reject(last(names))(newtable)
end

reapply(transform::Remainder, table, cache) =
  _apply(transform, table, cache) |> first


using DifferentialGeometries
using ForwardDiff, LinearAlgebra, Random, Test

Random.seed!(0)

@testset "DifferentialGeometries tests" begin
  include("Cartesian.jl")
  include("Cylindrical.jl")
  include("Toroidal.jl")
  include("ChangeOfCoordinates.jl")
  include("AutodiffPropagation.jl")
end


#cuda examples
using CUDA, Test, CUDA.CUFFT
function pi_gpu(n)
    4 * sum( CURAND(Float64, n).^2 .+ CURAND(Float64,n).^2 <= 1) / n
end

b = copy(a)
fill!(b, 0)
@btime b == CUDA.zeros(Int, 1024)

a = CUDA.rand(2,2)
map(sin, a)
@show fft(a)


"""
A `DiscreteDictionary{I,T}` is a dictionary where `I` is a discrete type, e.g.,
the integers. Expansions in this discrete set typically constitute a vector space.

The elements of the set are the Euclidean unit vectors, with a finite index set.
"""
abstract type DiscreteDictionary{I,T} <: Dictionary{I,T}
end

coefficienttype(::Type{DiscreteDictionary{I,T}}) where {I,T} = T

# The point x is in the support of d exactly when it is within the bounds of
# the index set, so we can do a checkbounds with Bool argument (which does not
# throw an error but returns true or false).
dict_in_support(d::DiscreteDictionary{I}, idx, x::I) where I = checkbounds(Bool, d, x)

function dict_in_support(d::DiscreteDictionary{I}, idx, x) where I
    try
        # Attempt to convert x to a native index. If that fails, the result is false.
        idxn = native_index(d, x)
        checkbounds(Bool, d, idxn)
    catch e
        false
    end
end

# Evaluation of discrete sets works as follows:
# -> eval_element: does bounds check on idx
#   -> unsafe_eval_element1: in_support corresponds to bounds check on x
#     -> unsafe_eval_element: return zero(T), or one(T) if idx == x
unsafe_eval_element(set::DiscreteDictionary{I,T}, idx::I, x::I) where {I,T} =
    idx == x ? one(T) : zero(T)

# In the routine above, we made sure that idx and x have the same type, so that
# they can be compared. If they do not have the same type, we can try to
# promote the indices.
unsafe_eval_element(set::DiscreteDictionary, idx, x) =
    unsafe_eval_element(set, native_index(set, idx), native_index(set, x))

isdiscrete(dict::Dictionary) = false
isdiscrete(dict::DiscreteDictionary) = true

#########################
# Concrete discrete sets
#########################

# We have to make these concrete sets because they store the size, and this
# depends on the type of the data. This is so in spite of the fact that type
# parameter I (the index type) already contains information.

"""
A `DiscreteVectorDictionary{T}` describes the linear space of vectors of finite length
`n` with element type `T`.
"""
struct DiscreteVectorDictionary{T} <: DiscreteDictionary{Int,T}
    n   ::  Int
end

# We set a default codomain type Float64
DiscreteVectorDictionary(n::Int) = DiscreteVectorDictionary{Float64}(n)

size(d::DiscreteVectorDictionary) = (d.n,)

similar(d::DiscreteVectorDictionary{T}, ::Type{Int}, n::Int) where {T} = DiscreteVectorDictionary{T}(n)

support(d::DiscreteVectorDictionary) = ClosedInterval{Int}(1, length(d))

"""
A `DiscreteArrayDictionary{N,T}` describes the linear space of arrays of finite size
`size(d)` with element type `T`.
"""
struct DiscreteArrayDictionary{T,N} <: DiscreteDictionary{ProductIndex{N},T}
    size    ::  NTuple{N,Int}
end

DiscreteArrayDictionary(size...) = DiscreteArrayDictionary{Float64}(size...)

DiscreteArrayDictionary{T}(size::Int...) where {T} = DicreteArrayDictionary{T}(size)

DiscreteArrayDictionary{T}(size::NTuple{N,Int}) where {T,N} =
    DiscreteArrayDictionary{T,N}(size)

DiscreteArrayDictionary(array::AbstractArray{T,N}) where {T,N} =
    DiscreteArrayDictionary{T,N}(size(array))

size(d::DiscreteArrayDictionary) = d.size

similar(d::DiscreteArrayDictionary{T}, ::Type{ProductIndex{N}}, dims::Vararg{Int,N}) where {T,N} = DiscreteArrayDictionary{T}(dims)

ordering(d::DiscreteArrayDictionary{N}) where {N} = ProductIndexList{N}(size(d))

native_index(d::DiscreteArrayDictionary, idx) = product_native_index(size(d), idx)

"""
A `DiscreteMultiArrayDictionary{A,T}` describes the linear space of multi-arrays
with element type `T`.
"""
struct DiscreteMultiArrayDictionary{T} <: DiscreteDictionary{MultilinearIndex,T}
    offsets ::  Vector{Int}
end

size(d::DiscreteMultiArrayDictionary) = (d.offsets[end],)

unsafe_offsets(d::DiscreteMultiArrayDictionary) = d.offsets

function similar(d::DiscreteMultiArrayDictionary, T::Type{MultilinearIndex}, size::Int...)
    @assert length(d) == prod(size)
    similar(d, T, d.offsets)
end

similar(d::DiscreteMultiArrayDictionary{T}, ::Type{MultilinearIndex}, offsets::Vector{Int}) where {T} = DiscreteMultiArrayDictionary{T}(offsets)

ordering(d::DiscreteMultiArrayDictionary) = MultilinearIndexList(unsafe_offsets(d))

support(d::DiscreteMultiArrayDictionary) = ClosedInterval{Int}(1, length(d))


type VineGraphChart <: JFinM_Charts.AbstractChart
    chartType::String
    cmd::String
    shape::String
    maxLayer::Int
    styles::Array{String, 1}
    colors::Array{String, 1}
    lens::Array{String, 1}
    weights::Array{String, 1}
end

## define default values
VineGraphChart() = VineGraphChart("VineGraphChart",
                                  "neato",
                                  "circle",
                                  2,
                                  ["bold", "dotted"],
                                  ["black", "black"],
                                  ["1.0", "1.8"],
                                  ["1", "1"]
                                  )

function GViz(vn::Copulas.Vine,
              chrt::VineGraphChart)
    nVars = size(vn.trees, 1)
    layers = Copulas.linkLayers(vn)
    dotCode = ""
    dotCode = string(dotCode, "graph G {\n")
    dotCode = string(dotCode, "node [shape=$(chrt.shape)];\n")
    for ii=1:nVars
        for jj=ii+1:nVars
            currLayer = layers[ii, jj]
            if currLayer <= chrt.maxLayer
                tmpCode = "$ii -- $jj [style=$(chrt.styles[currLayer]); color=$(chrt.colors[currLayer]); len=$(chrt.lens[currLayer]); weight=$(chrt.weights[currLayer])];\n"
                    dotCode = string(dotCode, tmpCode)
            end
        end
    end
    dotCode = string(dotCode, "}")
    return GViz("VineGraphChart", "neato", dotCode)
end


using Compat
using Compat.Printf

include("BoxFMMType.jl")
include("BoxFMMUtil.jl")
include("MBHFMM2D.jl")

using PyPlot

## define source and target points

ns = 100
nt = 20

# sources on an ellipse, targets inside

src = Array{Float64}(undef,2,ns)
srcsort = copy(src)
isrcsort = Array{Int32}(undef,ns)
targ = Array{Float64}(undef,2,nt)
targsort = copy(targ)
itargsort = Array{Int32}(undef,nt)

center = [6.1;-1.3]

h = 2.0*pi/ns
a = 1.123
b = 0.7

for i = 1:ns
    t = (i-1)*h
    src[:,i] = center + [a*cos(t);b*sin(t)]
end

for i = 1:nt
    xt = randn(); yt = randn();
    elnorm = sqrt((xt/a)^2 + (yt/b)^2)
    sc = rand();
    xt = sc*xt/elnorm; yt = sc*yt/elnorm
    targ[:,i] = center + [xt;yt]
end

## build tree based on points

maxboxes = 100000 # length of arrays allocated

# no more than 30 sources and targets per box
tree, sorted_pts, ier = BoxTree2DMaxBoxesST(src,targ,maxboxes,
                                 maxnodes=30)

@printf "built tree with %d boxes and %d levels\n" tree.nboxes tree.nlev

println("sorted_pts ",sorted_pts.srcsort)

@printf "plotting tree and points ...\n"

boxfmm2d_plot(tree)

plot(src[1,:],src[2,:],"bo")
plot(targ[1,:],targ[2,:],"rx")


#Illustrate how to set up certain constraint sets that are not found in other examples

using SetIntersectionProjection
using MAT
using LinearAlgebra, SparseArrays
using Statistics
using PyPlot
using JOLI

#define computational grid structure
mutable struct compgrid
  d :: Tuple
  n :: Tuple
end

#select working precision
TF=Float32 #Float32

#load escalator video) (Mablab files for compatibility with matlab only solvers for comparison...)
if isfile("escalator_data.mat") == true
  file     = matopen("escalator_data.mat")
  #file = matopen(joinpath(dirname(pathof(SetIntersectionProjection)), "../examples/Data/escalator_data.mat"))
else
  println("downloading escalator video from http://cvxr.com/tfocs/demos/rpca/escalator_data.mat")
  run(`wget http://cvxr.com/tfocs/demos/rpca/escalator_data.mat`)
  file     = matopen("escalator_data.mat")
  #error("download escalator video from http://cvxr.com/tfocs/demos/rpca/escalator_data.mat")
end

mtrue    = read(file, "X");
n1       = convert(Integer,read(file, "m"));
n2       = convert(Integer,read(file, "n"));
m_mat    = convert(Array{TF,2},mtrue);
m_tensor = convert(Array{TF,3},reshape(mtrue,n1,n2,Integer(200)));

#computational grid for the video
comp_grid  = compgrid((1f0,1f0,1f0),(size(m_tensor,1),size(m_tensor,2), size(m_tensor,3)));

comp_grid_time_slice = compgrid((1f0,1f0),(size(m_tensor,1),size(m_tensor,2)));

######################################################################
######## (anisotropic) total-variation on the time-coordinate ########

#the video has coordinates x-y-time, so the time corresponds to the 3rd coordinate
#which is z in the SetIntersectionProjection system (x-y-z)

#initialize constraints
constraint = Vector{SetIntersectionProjection.set_definitions}()
options = PARSDMM_options()

#l1 (total variation) constraints (in one direction)

#find a reasonable value for the l1-ball
(TD_OP, AtA_diag, dense, TD_n) = get_TD_operator(comp_grid,"D_z",options.FL);
TV_z = norm(TD_OP*vec(m_tensor),1)

m_min     = 0.0
m_max     = 0.025*TV_z
set_type  = "l1"
TD_OP     = "D_z"
app_mode  = ("tensor","")
custom_TD_OP = ([],false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP));

(P_sub,TD_OP,set_Prop) = setup_constraints(constraint,comp_grid,options.FL);
(TD_OP,AtA,l,y)        = PARSDMM_precompute_distribute(TD_OP,set_Prop,comp_grid,options);

@time (x,log_PARSDMM) = PARSDMM(vec(m_tensor),AtA,TD_OP,set_Prop,P_sub,comp_grid,options);
@time (x,log_PARSDMM) = PARSDMM(vec(m_tensor),AtA,TD_OP,set_Prop,P_sub,comp_grid,options);

m_proj = reshape(x,comp_grid.n);

figure();
for i=1:comp_grid.n[3]
    subplot(2,1,1);imshow(m_proj[:,:,i],cmap="gray");title("projected");
    subplot(2,1,2);imshow(m_tensor[:,:,i],cmap="gray");title("original");
    pause(0.01)
end

###############################################################################
######## (anisotropic) total-variation on the time-derivative          ########
########                     using JOLI operators                      ########

options = PARSDMM_options();

#TV operator per time-slice
(TV, AtA_diag, dense, TD_n) = get_TD_operator(comp_grid_time_slice,"TV",options.FL);
#time derivative over the time-slices
(D, AtA_diag, dense, TD_n)  =  get_TD_operator(comp_grid,"D_z",options.FL);

CustomOP_explicit_sparse = kron(TV, SparseMatrixCSC{TF}(LinearAlgebra.I, comp_grid.n[3]-1,comp_grid.n[3]-1))* D;

D  = joMatrix(D);
TV = joMatrix(TV);
CustomOP_JOLI = joKron(TV, joEye(comp_grid.n[3]-1,DDT=Float32,RDT=Float32))* D;

## Solve using JOLI ##

    #initialize constraints
    constraint = Vector{SetIntersectionProjection.set_definitions}();

    m_min     = 0.0
    m_max     = 0.1*norm(CustomOP_JOLI*vec(m_tensor),1)
    set_type  = "l1"
    TD_OP     = "identity"
    app_mode  = ("matrix","")
    custom_TD_OP = (CustomOP_JOLI,false)
    push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP));

    (P_sub,TD_OP,set_Prop) = setup_constraints(constraint,comp_grid,options.FL);

    #set properties of custom operator
    set_Prop.AtA_diag[1]  = false
    set_Prop.dense[1]     = false
    set_Prop.banded[1]    = true

    (TD_OP,AtA,l,y)        = PARSDMM_precompute_distribute(TD_OP,set_Prop,comp_grid,options);

    @time (x_joli,log_PARSDMM) = PARSDMM(vec(m_tensor),AtA,TD_OP,set_Prop,P_sub,comp_grid,options);
    @time (x_joli,log_PARSDMM) = PARSDMM(vec(m_tensor),AtA,TD_OP,set_Prop,P_sub,comp_grid,options);

## solve using explicit sparse array ##
    constraint = Vector{SetIntersectionProjection.set_definitions}();

    m_min     = 0.0
    m_max     = 0.025*norm(CustomOP_explicit_sparse*vec(m_tensor),1)
    set_type  = "l1"
    TD_OP     = "identity"
    app_mode  = ("matrix","")
    custom_TD_OP = (CustomOP_explicit_sparse,false)
    push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP));

    options=PARSDMM_options()
    (P_sub,TD_OP,set_Prop) = setup_constraints(constraint,comp_grid,options.FL);

    #set properties of custom operator
    set_Prop.AtA_diag[1]  = false
    set_Prop.dense[1]     = false
    set_Prop.banded[1]    = true

    (TD_OP,AtA,l,y)        = PARSDMM_precompute_distribute(TD_OP,set_Prop,comp_grid,options);

    @time (x_sp,log_PARSDMM) = PARSDMM(vec(m_tensor),AtA,TD_OP,set_Prop,P_sub,comp_grid,options);
    @time (x_sp,log_PARSDMM) = PARSDMM(vec(m_tensor),AtA,TD_OP,set_Prop,P_sub,comp_grid,options);

    m_proj = reshape(x_sp,comp_grid.n)

figure();
for i=1:comp_grid.n[3]
    subplot(2,1,1);imshow(m_proj[:,:,i],cmap="gray");title("projected");
    subplot(2,1,2);imshow(m_tensor[:,:,i],cmap="gray");title("original");
    pause(0.025)
end

###############################################################################
######## (anisotropic) total-variation on every subtracted time-slice #########

# every  time-slice should have a small TV like this: || TV(slice_2 -slice_1) ||_1 <= sigma, || TV(slice_3 -slice_1) ||_1 <= sigma, ...

#initialize constraints
# constraint = Vector{SetIntersectionProjection.set_definitions}()

# slice_ref = copy(m_tensor[:,:,1])

# m_shifted = similar(m_tensor)
# for i=1:comp_grid.n[3]
#     m_shifted[:,:,i] = m_tensor[:,:,i] - slice_ref
# end

# #find a reasonable value for the l1-ball
# (TD_OP, AtA_diag, dense, TD_n) = get_TD_operator(comp_grid,"D_z",options.FL)
# TV_z = norm(TD_OP*vec(m_shifted),1)

# m_min     = 0.0
# m_max     = 0.1*TV_z
# set_type  = "l1"
# TD_OP     = "TV"
# app_mode  = ("matrix","")
# custom_TD_OP = ([],false)
# push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

# options=PARSDMM_options()
# (P_sub,TD_OP,set_Prop) = setup_constraints(constraint,comp_grid,options.FL)
# (TD_OP,AtA,l,y)        = PARSDMM_precompute_distribute(TD_OP,set_Prop,comp_grid,options)

# @time (x,log_PARSDMM) = PARSDMM(vec(m_shifted),AtA,TD_OP,set_Prop,P_sub,comp_grid,options);
# @time (x,log_PARSDMM) = PARSDMM(vec(m_shifted),AtA,TD_OP,set_Prop,P_sub,comp_grid,options);

# m_proj = reshape(x,comp_grid.n)
# for i=1:comp_grid.n[3]
#     m_proj[:,:,i] = m_proj[:,:,i] + slice_ref
# end

# figure();
# for i=1:comp_grid.n[3]
#     subplot(2,1,1);imshow(m_proj[:,:,i],cmap="gray");title("projected");
#     subplot(2,1,2);imshow(m_tensor[:,:,i],cmap="gray");title("original");
#     pause(0.025)
# end

###############################################################################
########################## custom dense JOLI matrix  ##########################

x = randn(Float32,10,20)
comp_grid  = compgrid((1f0,1f0),(10,20))

constraint = Vector{SetIntersectionProjection.set_definitions}()

options=PARSDMM_options()

#l1:
m_min     = 0.0f0
m_max     = 10.0f0
set_type  = "l1"
TD_OP     = "identity"
app_mode  = ("matrix","")
custom_TD_OP = (joMatrix(randn(options.FL, 133,prod(comp_grid.n))),false)
push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

(P_sub,TD_OP,set_Prop) = setup_constraints(constraint,comp_grid,options.FL)

#set properties of custom operator
set_Prop.AtA_diag[1]  = false
set_Prop.dense[1]     = true
set_Prop.banded[1]    = false

options.adjust_rho             = false
options.adjust_feasibility_rho = false

(TD_OP,AtA,l,y)        = PARSDMM_precompute_distribute(TD_OP,set_Prop,comp_grid,options)

@time (x_proj,log_PARSDMM) = PARSDMM(vec(x),AtA,TD_OP,set_Prop,P_sub,comp_grid,options);
@time (x_proj,log_PARSDMM) = PARSDMM(vec(x),AtA,TD_OP,set_Prop,P_sub,comp_grid,options);


using Test

@test 1+1 == 2


function makeMesh(r,outDir=".")

  n = size(r,1)

  gmsh.initialize()
  gmsh.option.setNumber("General.Terminal", 0) # change to 1 for Gmsh output

  gmsh.model.add("mixing")

  lc = 1e-2
  global theta = range(0.0,stop=2.0*pi,length=n+1)

  for i=1:n
    gmsh.model.geo.addPoint(r[i]*cos(theta[i]),
                            r[i]*sin(theta[i]),
                            0.0, lc, i)
  end

  for i=1:n
    gmsh.model.geo.addPoint(2*cos(theta[i]),2*sin(theta[i]),0.0, 2*lc,n+i)
  end

  nCircle = 15
  global theta1 = range(0.0,stop=2.0*pi,length=nCircle+1)
  for i=1:nCircle
    gmsh.model.geo.addPoint(1.25+0.1*cos(theta1[i]),1.25+0.1*sin(theta1[i]),0.0, lc,2n+i)
#    gmsh.model.geo.addPoint(1.75+0.1*cos(theta1[i]),0.0+0.1*sin(theta1[i]),0.0, lc,2n+i)
  end

  #  Assemble boundary points into curves following the right-hand-rule
  #  (traverse the boundary with your right hand toward the "outside")
  gmsh.model.geo.addBSpline([1;collect(n:-1:1)],1)
  gmsh.model.geo.addSpline([collect((n+1):(2*n));n+1],2)
  gmsh.model.geo.addSpline([2*n+1;collect(2*n+nCircle:-1:2*n+1)],3)

  gmsh.model.geo.addCurveLoop([1],1)
  gmsh.model.geo.addCurveLoop([2],2)
  gmsh.model.geo.addCurveLoop([3],3)
  gmsh.model.geo.addPlaneSurface([1,2,3],1)
  gmsh.model.geo.addPlaneSurface([3],2)

  gmsh.model.addPhysicalGroup(1, [1], 1)
  gmsh.model.addPhysicalGroup(1, [2], 2)
 # gmsh.model.addPhysicalGroup(1, [3], 3)
  gmsh.model.setPhysicalName(1, 1, "Inner")
  gmsh.model.setPhysicalName(1, 2, "Outer")

  gmsh.model.addPhysicalGroup(2, [1], 1)
  gmsh.model.addPhysicalGroup(2, [2], 2)
  gmsh.model.setPhysicalName(2, 1, "Mesh")
  # gmsh.model.setPhysicalName(2, 2, "Sensor")

  gmsh.model.geo.synchronize()

  gmsh.model.mesh.generate(2)
  gmsh.model.mesh.setOrder(2)  # quadratic elements

  ##nodeIds,nodeCoords,_ = gmsh.model.mesh.getNodes()
  ##x = reshape(nodeCoords,3,:)
  nodeTags1,xx1 = gmsh.model.mesh.getNodesForPhysicalGroup(2,1)
  nodeTags2,xx2 = gmsh.model.mesh.getNodesForPhysicalGroup(2,2)
  xx1 = reshape(xx1,3,:)
  xx2 = reshape(xx2,3,:)
  nn = convert(Int64,maximum(nodeTags2))
  x = zeros(3,nn)
  for i=1:length(nodeTags1)
    x[:,nodeTags1[i]] = xx1[:,i]
  end
  for i=1:length(nodeTags2)
    x[:,nodeTags2[i]] = xx2[:,i]
  end

  nInner,xInner = gmsh.model.mesh.getNodesForPhysicalGroup(1,1) # 1D groups, #1 is the inner circle
  nOuter,xOuter = gmsh.model.mesh.getNodesForPhysicalGroup(1,2) # 1D groups, #2 is the outer circle

  elemTypes1, elemTags1, elemConnectivity1 = gmsh.model.mesh.getElements(2,1)
  elemTypes2, elemTags2, elemConnectivity2 = gmsh.model.mesh.getElements(2,2)

  eConn1 = reshape(elemConnectivity1[1],6,:)
  eConn2 = reshape(elemConnectivity2[1],6,:)

  eConn = [eConn1 eConn2]
#  eC = convert(Array{Int64,2},eConn)   # for debugging

#  gmsh.write(outDir*"mixing.msh")
  gmsh.finalize()

 return x,eConn,eConn2, nInner,xInner, nOuter,xOuter
end


using BenchmarkTools: @btime

function memory(start, target)
  numbers = zeros(Int32, target)
  for (i,n) in enumerate(start[1:end-1])
    numbers[n+1]=i
  end
  turn = length(start)
  number = start[end]
  while turn < target
    seen = numbers[number+1]
    numbers[number+1] = turn
    number = seen == 0 ? 0 : turn - seen
    turn += 1
  end
  number
end

input = [0,5,4,1,10,14,7]
part1 = memory(input, 2020)
println("$part1")

part2 = @btime memory(input, 30000000)
println("$part2")


"Utility functions"

import ForneyLab: cholinv
import SpecialFunctions: digamma

function wMatrix(γ, order)
    mW = 1e8*Matrix{Float64}(I, order, order)
    mW[end, end] = γ
    return mW
end

# function cholinv(M::AbstractMatrix)
#     return LinearAlgebra.inv(M)
# end

function mvdigamma(x, order)
    "Multivariate digamma function (see https://en.wikipedia.org/wiki/Multivariate_gamma_function)"
    result = 0.
    for i = 1:order
        result += digamma(x + (1 - i)/2.)
    end
    return result
end


abstract type StepController end
struct ConstantSteps <: StepController end
function propose_step!(::ConstantSteps, integ)
    return integ.dt
end

struct StandardSteps <: StepController end
function propose_step!(::StandardSteps, integ)

    @unpack dt, EEst = integ
    @unpack q = integ.constants
    @unpack gamma, qmin, qmax = integ.opts

    if iszero(EEst)
        q = inv(qmax)
    else
        localconvrate = q+1
        qtmp = DiffEqBase.fastpow(EEst, 1/localconvrate) / gamma
        @fastmath q = DiffEqBase.value(max(inv(qmax), min(inv(qmin), qtmp)))
    end
    dt /= q
    return dt
end

struct PISteps <: StepController end
function propose_step!(::PISteps, integ)
    # PI-controller
    @unpack dt, EEst, qold = integ
    @unpack gamma, beta1, beta2, qmin, qmax = integ.opts

    if iszero(EEst)
        q = inv(qmax)
    elseif EEst < 1  # Accepted
        q11 = DiffEqBase.value(DiffEqBase.fastpow(EEst, beta1))
        q = q11 / DiffEqBase.fastpow(qold, beta2) / gamma

        @fastmath q = DiffEqBase.value(max(inv(qmax),min(inv(qmin), q)))
        integ.qold = max(integ.EEst, integ.opts.qoldinit)
    else # Rejected step
        q11 = DiffEqBase.value(DiffEqBase.fastpow(EEst, beta1))
        q = q11 / gamma
        @fastmath q = DiffEqBase.value(max(inv(qmax),min(inv(qmin), q)))
    end
    dt /= q
    return dt
end


module DoubleDouble

export DoubleF64, ComplexDF64

import Printf
import Base: +, -, *, /, ^, <, ==, <=

# Basic building blocks
"""
    quick_two_sum(a, b)

Computes `s = fl(a+b)` and `e = err(a+b)`. Assumes `|a| ≥ |b|`.
"""
@inline function quick_two_sum(a::Float64, b::Float64)
    s = a + b
    e = b - (s - a)
    s, e
end

"""
    two_sum(a, b)

Computes `s = fl(a+b)` and `e = err(a+b)`.
"""
@inline function two_sum(a::Float64, b::Float64)
    s = a + b
    v = s - a
    e = (a - (s - v)) + (b - v)

    s, e
end

"""
    two_diff(a, b)

Computes `s = fl(a-b)` and `e = err(a-b)`.
"""
@inline function two_diff(a::Float64, b::Float64)
    s = a - b
    v = s - a
    e = (a - (s - v)) - (b + v)

    s, e
end

"""
    two_prod(a, b)

Computes `s = fl(a*b)` and `e = err(a*b)`.
"""
@inline function two_prod(a::Float64, b::Float64)
    p = a * b
    e = fma(a, b, -p)
    p, e
end

"""
    two_square(a)

Computes `s = fl(a*a)` and `e = err(a*a)`. Faster than [`two_prod(a, a)`](@ref).
"""
@inline function two_square(a::Float64)
    p = a * a
    e = fma(a, a, -p)
    p, e
end

"""
    DoubleF64(x [, mode::ComputeMode]) <: AbstractFloat
"""
struct DoubleF64 <: AbstractFloat
    hi::Float64
    lo::Float64
end

DoubleF64(x::DoubleF64) = DoubleF64(x.hi, x.lo)
function DoubleF64(x::Float64)
    DoubleF64(x, isinf(x) ? Inf : 0.0)
end

function DoubleF64(x::Float32)
    DoubleF64(convert(Float64, x), isinf(x) ? Inf : 0.0)
end
function DoubleF64(x::Float16)
    DoubleF64(convert(Float64, x), isinf(x) ? Inf : 0.0)
end
function DoubleF64(x::Integer)
    DoubleF64(convert(Float64, x), isinf(x) ? Inf : 0.0)
end
function DoubleF64(x::BigFloat)
    z = convert(Float64, x)
    DoubleF64(z, convert(Float64, x - z))
end
function DoubleF64(x::Irrational)
    DoubleF64(big(x))
end
function DoubleF64(x::Rational)
    DoubleF64(numerator(x)) / DoubleF64(denominator(x))
end

const ComplexDF64 = Complex{DoubleF64}

hi(x::DoubleF64) = x.hi
lo(x::DoubleF64) = x.lo

Base.isbits(::DoubleF64) = true
Base.isbitstype(::Type{DoubleF64}) = true

Base.zero(::DoubleF64) = DoubleF64(0.0, 0.0)
Base.zero(::Type{DoubleF64}) = DoubleF64(0.0, 0.0)

Base.one(::DoubleF64) = DoubleF64(1.0, 0.0)
Base.one(::Type{DoubleF64}) = DoubleF64(1.0, 0.0)

Base.convert(::Type{T}, a::DoubleF64) where {T<:AbstractFloat} = convert(T, a.hi)
Base.convert(::Type{BigFloat}, a::DoubleF64) = big(a.hi) + big(a.lo)
Base.BigFloat(a::DoubleF64) = convert(BigFloat, a)
Base.convert(::Type{T}, a::DoubleF64) where {T<:Integer} = convert(T, a.hi)
Base.convert(::Type{Integer}, a::DoubleF64) = convert(Int64, a.hi)
Base.convert(::Type{BigInt}, a::DoubleF64) = convert(BigInt, big(a.hi) + big(a.lo))

Base.convert(::Type{DoubleF64}, x::DoubleF64) = x
Base.convert(::Type{DoubleF64}, x::AbstractFloat) = DoubleF64(x)
Base.convert(::Type{DoubleF64}, x::Irrational) = DoubleF64(x)
Base.convert(::Type{DoubleF64}, x::Integer) = DoubleF64(x)

Base.promote_rule(::Type{DoubleF64}, ::Type{<:Integer}) = DoubleF64
Base.promote_rule(::Type{DoubleF64}, ::Type{BigInt}) = BigFloat
Base.promote_rule(::Type{DoubleF64}, ::Type{BigFloat}) = BigFloat
Base.promote_rule(::Type{DoubleF64}, ::Type{Float64}) = DoubleF64
Base.promote_rule(::Type{DoubleF64}, ::Type{Float32}) = DoubleF64
Base.promote_rule(::Type{DoubleF64}, ::Type{Float16}) = DoubleF64

Base.big(x::DoubleF64) = big(x.hi) + big(x.lo)

const double_π = DoubleF64(pi)
const double_2pi = DoubleF64(2 * BigFloat(Base.pi))
const double_pi = DoubleF64(BigFloat(Base.pi))
const double_pi2 = DoubleF64(BigFloat(Base.pi) * 0.5)
const double_pi4 = DoubleF64(BigFloat(Base.pi) * 0.25)
const double_pi16 = DoubleF64(BigFloat(Base.pi) * (1 / 16))
const double_3pi4 = DoubleF64(BigFloat(Base.pi) * 0.75)
const double_nan = DoubleF64(NaN, NaN)
const double_inf = DoubleF64(Inf)
const double_e = convert(DoubleF64, Base.MathConstants.e)

const double_log2 = convert(DoubleF64, log(BigFloat(2.0)))
const double_log10 = convert(DoubleF64, log(BigFloat(10.0)))

const double_eps = 4.93038065763132e-32 # 2^-104

#
# ADDITION
#
"""
    wide_add(a::Float64, b::Float64)

Add two `Float64`s with `DoubleF64` precision.
"""
@inline function wide_add(a::Float64, b::Float64)
    hi, lo = two_sum(a, b)

    DoubleF64(hi, lo)
end

@inline function +(a::DoubleF64, b::Float64)
    hi, lo = two_sum(a.hi, b)
    lo += a.lo
    hi, lo = quick_two_sum(hi, lo)

    DoubleF64(hi, lo)
end
+(a::Float64, b::DoubleF64) = b + a
+(a::Integer, b::DoubleF64) = b + float(a)
+(a::DoubleF64, b::Integer) = a + float(b)

@inline function +(a::DoubleF64, b::DoubleF64)
    hi, lo = two_sum(a.hi, b.hi)
    lo += (a.lo + b.lo)
    hi, lo = quick_two_sum(hi, lo)

    DoubleF64(hi, lo)
end

#
# SUBSTRACTION
#
"""
    wide_sub(a::Float64, b::Float64)

Subtract two `Float64`s with `DoubleF64` precision.
"""
@inline function wide_sub(a::Float64, b::Float64)
    hi, lo = two_diff(a, b)

    DoubleF64(hi, lo)
end

@inline function -(a::DoubleF64, b::Float64)
    hi, lo = two_diff(a.hi, b)
    lo += a.lo
    hi, lo = quick_two_sum(hi, lo)

    DoubleF64(hi, lo)
end

@inline function -(a::Float64, b::DoubleF64)
    hi, lo = two_diff(a, b.hi)
    lo -= b.lo
    hi, lo = quick_two_sum(hi, lo)

    DoubleF64(hi, lo)
end

@inline function -(a::DoubleF64, b::DoubleF64)
    hi, lo = two_diff(a.hi, b.hi)
    lo += a.lo
    lo -= b.lo
    hi, lo = quick_two_sum(hi, lo)

    DoubleF64(hi, lo)
end
-(a::DoubleF64) = DoubleF64(-a.hi, -a.lo)

#
# MULTIPLICATION
#
"""
    wide_mutiply(a::Float64, b::Float64)

Multiply two `Float64`s with `DoubleF64` precision.
"""
@inline function wide_mul(a::Float64, b::Float64)
    hi, lo = two_prod(a, b)
    DoubleF64(hi, lo)
end

@inline function *(a::DoubleF64, b::Float64)
    p1, p2 = two_prod(a.hi, b)
    p2 += a.lo * b
    p1, p2 = quick_two_sum(p1, p2)

    DoubleF64(p1, p2)
end
*(a::Float64, b::DoubleF64) = b * a
*(a::Integer, b::DoubleF64) = b * float(a)
*(a::Bool, b::DoubleF64) = b * convert(Float64, a)
*(a::DoubleF64, b::Integer) = a * float(b)

@inline function *(a::DoubleF64, b::DoubleF64)
    p1, p2 = two_prod(a.hi, b.hi)
    p2 += a.hi * b.lo + a.lo * b.hi
    p1, p2 = quick_two_sum(p1, p2)

    DoubleF64(p1, p2)
end

#
# DIVISION
#
"""
    wide_div(a::Float64, b::Float64)

Divide two `Float64`s with `DoubleF64` precision.
"""
@inline function wide_div(a::Float64, b::Float64)
    q1 = a / b
    # Compute a - q1 * b
    p1, p2 = two_prod(q1, b)
    s, e = two_diff(a, p1)
    e -= p2

    # get next approximation
    q2 = (s + e) / b

    s, e = quick_two_sum(q1, q2)

    DoubleF64(s, e)
end

@inline function /(a::DoubleF64, b::Float64)
    q1 = a.hi / b

    # Compute  this - q1 * d
    p1, p2 = two_prod(q1, b)
    s, e = two_diff(a.hi, p1)
    e += a.lo
    e -= p2

    # get next approximation.
    q2 = (s + e) / b

    # renormalize
    hi, lo = quick_two_sum(q1, q2)

    DoubleF64(hi, lo)
end

@inline function /(a::DoubleF64, b::DoubleF64)
    q1 = a.hi / b.hi  # approximate quotient

    # compute  this - q1 * dd
    r = b * q1
    s1, s2 = two_diff(a.hi, r.hi)
    s2 -= r.lo
    s2 += a.lo

    # get next approximation
    q2 = (s1 + s2) / b.hi

    # renormalize
    hi, lo = quick_two_sum(q1, q2)

    DoubleF64(hi, lo)
end

/(a::Float64, b::DoubleF64) = DoubleF64(a) / b
/(a::Integer, b::DoubleF64) = float(a) / b
/(a::DoubleF64, b::Integer) = a / float(b)

Base.inv(a::DoubleF64) = 1.0 / a

Base.rem(a::DoubleF64, b::DoubleF64) = a - round(a / b) * b
@inline function Base.divrem(a::DoubleF64, b::DoubleF64)
    n = round(a / b)
    n, a - n * b
end

function Base.mod(x::DoubleF64, y::DoubleF64)
    n = round(a / b)
    return (a - b * n)
end

#
# POWERS
#

"""
    square(x::DoubleF64)

Compute `x * x` in a more efficient way.
"""
@inline function square(a::DoubleF64)
    p1, p2 = two_square(a.hi)
    p2 += 2.0 * a.hi * a.lo
    p2 += a.lo * a.lo
    hi, lo = quick_two_sum(p1, p2)
    return DoubleF64(hi, lo)
end
square(a::Float64) = wide_square(a)
"""
    wide_square(x::Float64)

Convert `x` to a DoubleF64 and then compute `x*x`.
"""
@inline function wide_square(a::Float64)
    hi, lo = two_square(a)
    return DoubleF64(hi, lo)
end
# Implementation adapted from Base
@inline function power_by_squaring(x::DoubleF64, p::Integer)
    if p == 1
        return copy(x)
    elseif p == 0
        return one(x)
    elseif p == 2
        return square(x)
    end
    P = abs(p)
    t = trailing_zeros(P) + 1
    P >>= t
    while (t -= 1) > 0
        x = square(x)
    end
    y = x
    while P > 0
        t = trailing_zeros(P) + 1
        P >>= t
        while (t -= 1) >= 0
            x = square(x)
        end
        y *= x
    end
    if (p < 0)
        return 1.0 / y
    end
    y
end

^(a::DoubleF64, p::Integer) = power_by_squaring(a, p)
^(a::DoubleF64, b::Real) = iszero(a) ? one(a) : exp(b * log(a))

@inline function Base.sqrt(a::DoubleF64)
    # Strategy:  Use Karp's trick:  if x is an approximation
    # to sqrt(a), then
    #
    #  sqrt(a) = a*x + [a - (a*x)^2] * x / 2   (approx)
    #
    # The approximation is accurate to twice the accuracy of x.
    # Also, the multiplication (a*x) and [-]*x can be done with
    # only half the precision.
    if a.hi < 0
        throw(DomainError("sqrt will only return a complex result if called with a complex argument."))
    end
    if iszero(a)
        return zero(a)
    end

    x = inv(sqrt(a.hi))
    ax = a.hi * x

    wide_add(ax, (a - square(ax)).hi * (x * 0.5))
end

"""
    wide_sqrt(x::Float64)

Convert `x` to a DoubleF64 and then compute `sqrt(x)`.
"""
wide_sqrt(a::Float64) = sqrt(DoubleF64(a))

# COMPARISON & EQUALITY
<(a::DoubleF64, b::DoubleF64) = a.hi + a.lo < b.hi + b.lo
<(a::DoubleF64, b::Float64) = a.hi < b || (a.hi == b) && a.lo < 0.0
<(a::Float64, b::DoubleF64) = a < b.hi || (a == b.hi) && b.lo > 0.0

Base.isless(a::DoubleF64, b::DoubleF64) = isless(a.hi + a.lo, b.hi + b.lo)

<=(a::DoubleF64, b::DoubleF64) = !(b < a)
<=(a::DoubleF64, b::Float64) = !(b < a)
<=(a::Float64, b::DoubleF64) = !(b < a)

==(a::DoubleF64, b::Float64) = a.hi == b && a.lo == 0.0
==(a::Float64, b::DoubleF64) = b == a
Base.iszero(a::DoubleF64) = a.hi == 0.0

Base.isone(a::DoubleF64) = a.hi == 1.0 && a.lo == 0.0
Base.abs(a::DoubleF64) = a.hi < 0.0 ? -a : a

Base.eps(::DoubleF64) = 4.93038065763132e-32 # 2^-104
Base.eps(::Type{DoubleF64}) = 4.93038065763132e-32 # 2^-104

# Base.realmin(::Type{DoubleF64}) = 2.0041683600089728e-292 # = 2^(-1022 + 53)
# Base.realmin(::Type{DoubleF64}) = 2.0041683600089728e-292 # = 2^(-1022 + 53)
# Base.realmax(::Type{DoubleF64}) = DoubleF64(1.79769313486231570815e+308, 9.97920154767359795037e+291);
# Base.realmax(::Type{DoubleF64}) = DoubleF64(1.79769313486231570815e+308, 9.97920154767359795037e+291);

Base.isnan(a::DoubleF64) = isnan(a.hi) || isnan(a.lo)
Base.isinf(a::DoubleF64) = isinf(a.hi)
Base.isfinite(a::DoubleF64) = isfinite(a.hi)

#
# ROUNDING
#
@inline function Base.round(a::DoubleF64, r::RoundingMode = RoundNearest)
    hi = round(a.hi, r)
    lo = 0.0

    if hi == a.hi
        # High word is an integer already.  Round the low word.
        lo = round(a.lo, r)

        # Renormalize. This is needed if hi = some integer, lo = 1/2.
        hi, lo = quick_two_sum(hi, lo)
    else
        # High word is not an integer.
        if abs(hi - a.hi) == 0.5 && a.lo < 0.0
            # There is a tie in the high word, consult the low word to break the tie.
            hi -= 1.0
        end
    end

    DoubleF64(hi, lo)
end

@inline function Base.floor(a::DoubleF64)
    hi = floor(a.hi)
    lo = 0.0

    if hi == a.hi
        lo = floor(a.lo)
        hi, lo = quick_two_sum(hi, lo)
    end

    DoubleF64(hi, lo)
end

@inline function Base.floor(::Type{I}, a::DoubleF64) where {I<:Integer}
    hi = floor(I, a.hi)
    lo = zero(I)

    if hi == a.hi
        lo = floor(I, a.lo)
    end

    hi + lo
end

@inline function Base.ceil(a::DoubleF64)
    hi = ceil(a.hi)
    lo = 0.0

    if hi == a.hi
        lo = ceil(a.lo)
        hi, lo = quick_two_sum(hi, lo)
    end

    DoubleF64(hi, lo)
end

@inline function Base.ceil(::Type{I}, a::DoubleF64) where {I<:Integer}
    hi = ceil(I, a.hi)
    lo = zero(I)

    if hi == a.hi
        lo = ceil(I, a.lo)
    end

    hi + lo
end

Base.trunc(a::DoubleF64) = a.hi ≥ 0.0 ? floor(a) : ceil(a)
Base.trunc(::Type{I}, a::DoubleF64) where {I<:Integer} =
    a.hi ≥ 0.0 ? floor(I, a) : ceil(I, a)
Base.isinteger(x::DoubleF64) = iszero(x - trunc(x))

# import Random: AbstractRNG, GLOBAL_RNG
# function Base.rand(rng::AbstractRNG, S::Type{DoubleF64})
#     u = rand(rng, UInt64)
#     f = Float64(u)
#     uf = UInt64(f)
#     ur = uf > u ? uf - u : u - uf
#     DoubleF64(5.421010862427522e-20 * f, 5.421010862427522e-20 * Float64(ur))
# end
# Base.rand(::Type{DoubleF64}) = rand(GLOBAL_RNG, DoubleF64)
#
# function Base.rand(rng::AbstractRNG, T::Type{<:DoubleF64}, dims::Vararg{Int, N}) where N
#     rands = Array{T}(dims)
#     for l in eachindex(rands)
#         rands[l] = rand(rng, T)
#     end
#     rands
# end
# Base.rand(T::Type{<:DoubleF64}, dims::Vararg{Int, N}) where N = rand(GLOBAL_RNG, T, dims)
#
# Base.rand(::Type{Complex{DoubleF64}}) = rand(Complex{DoubleF64})
# Base.rand(::Type{Complex{DoubleF64}}, dims::Vararg{Int, N}) where N = rand(Complex{DoubleF64}, dims)
# Base.rand(rng::AbstractRNG, ::Type{Complex{DoubleF64}}) = rand(rng, Complex{DoubleF64})
# Base.rand(rng::AbstractRNG, ::Type{Complex{DoubleF64}}, dims::Vararg{Int, N}) where N = rand(rng, Complex{DoubleF64}, dims)

function Base.decompose(a::DoubleF64)::Tuple{Int128,Int,Int}
    hi, lo = a.hi, a.lo
    num1, pow1, den1 = Base.decompose(hi)
    num2, pow2, den2 = Base.decompose(lo)

    num = Int128(num1)

    pdiff = pow1 - pow2
    shift = min(pdiff, 52)
    signed_num = den1 * (Int128(num) << shift) # den1 is +1/-1
    signed_num += den2 * (num2 >> (pdiff - shift)) # den2 is +1/-1

    num = abs(signed_num)
    den = signed_num ≥ 0 ? 1 : -1
    pow = pow1 - shift

    num, pow, den
end

# Precomputed valuee
const inv_fact = [DoubleF64(1.0 / BigFloat(factorial(Int64(k)))) for k = 3:17]
const ninv_fact = length(inv_fact)

Base.ldexp(a::DoubleF64, exp::Int) = DoubleF64(ldexp(a.hi, exp), ldexp(a.lo, exp))
"""
	mul_pwr2(a::DoubleF64, b::Float64)

`a` * `b`,  where `b` is a power of 2.
"""
mul_pwr2(a::DoubleF64, b::Float64) = DoubleF64(a.hi * b, a.lo * b)

function Base.exp(a::DoubleF64)
    # Strategy:  We first reduce the size of x by noting that
    #
    #       exp(kr + m * log(2)) = 2^m * exp(r)^k
    #
    # where m and k are integers.  By choosing m appropriately
    # we can make |kr| <= log(2) / 2 = 0.347.  Then exp(r) is
    # evaluated using the familiar Taylor series.  Reducing the
    # argument substantially speeds up the convergence.

    k = 512.0
    inv_k = 1.0 / k

    if a.hi ≤ -709.0
        return zero(a)
    end

    if a.hi ≥ 709.0
        convert(DoubleF64, Inf)
    end

    if iszero(a)
        return one(a)
    end

    if isone(a)
        return double_e
    end

    m = floor(a.hi / double_log2.hi + 0.5)
    r = mul_pwr2(a - double_log2 * m, inv_k)

    #
    p = square(r)
    s = r + mul_pwr2(p, 0.5)
    p *= r
    t = p * inv_fact[1]

    i = 1
    while true
        s += t
        p *= r
        i += 1
        t = p * inv_fact[i]

        if abs(convert(Float64, t)) < inv_k * double_eps || i > 6
            break
        end
    end

    s += t

    s = mul_pwr2(s, 2.0) + square(s)
    s = mul_pwr2(s, 2.0) + square(s)
    s = mul_pwr2(s, 2.0) + square(s)
    s = mul_pwr2(s, 2.0) + square(s)
    s = mul_pwr2(s, 2.0) + square(s)
    s = mul_pwr2(s, 2.0) + square(s)
    s = mul_pwr2(s, 2.0) + square(s)
    s = mul_pwr2(s, 2.0) + square(s)
    s = mul_pwr2(s, 2.0) + square(s)
    s += 1.0

    return ldexp(s, convert(Int, m))
end

Base.log10(a::DoubleF64) = log(a) / double_log10
Base.log2(a::DoubleF64) = log(a) / double_log2
function Base.log(a::DoubleF64)
    # Strategy.  The Taylor series for log converges much more
    # slowly than that of exp, due to the lack of the factorial
    # term in the denominator.  Hence this routine instead tries
    # to determine the root of the function
    #
    #     f(x) = exp(x) - a
    #
    # using Newton iteration.  The iteration is given by
    #
    #     x' = x - f(x)/f'(x)
    #        = x - (1 - a * exp(-x))
    #        = x + a * exp(-x) - 1.
    #
    # Only one iteration is needed, since Newton's iteration
    # approximately doubles the number of digits per iteration.

    if isone(a)
        return zero(DoubleF64)
    end

    if a.hi ≤ 0.0
        throw(DomainError("Passed non-positive argument to log."))
    end

    x = DoubleF64(log(a.hi)) # Initial approximation

    x + a * exp(-x) - 1.0
end

const sin_table = [DoubleF64(sin(k * big(π) * 0.0625)) for k = 1:4]
const cos_table = [DoubleF64(cos(k * big(π) * 0.0625)) for k = 1:4]

"""
	sin_taylor(a::DoubleF64)

Computes sin(a) using Taylor series. Assumes |a| <= π/32.
"""
function sin_taylor(a::DoubleF64)
    thresh = 0.5 * abs(convert(Float64, a)) * double_eps

    if iszero(a)
        return zero(a)
    end

    i = 1
    x = -square(a)
    s = a
    r = a
    while true
        r *= x
        t = r * inv_fact[i]
        s += t
        i += 2
        if i > ninv_fact || abs(convert(Float64, t)) < thresh
            break
        end
    end

    s
end

"""
	cos_taylor(a::DoubleF64)

Computes cos(a) using Taylor series. Assumes |a| <= π/32.
"""
function cos_taylor(a::DoubleF64)
    thresh = 0.5 * double_eps

    if iszero(a)
        return one(a)
    end

    x = -square(a)
    r = x
    s = 1.0 + mul_pwr2(r, 0.5)
    i = 2
    while true
        r *= x
        t = r * inv_fact[i]
        s += t
        i += 2
        if i > ninv_fact || abs(convert(Float64, t)) < thresh
            break
        end
    end

    s
end

function sincos_taylor(a::DoubleF64)
    if iszero(a)
        zero(a), one(a)
    end

    sin_a = sin_taylor(a)
    cos_a = sqrt(1.0 - square(sin_a))

    sin_a, cos_a
end

function Base.sin(a::DoubleF64)
    # Strategy.  To compute sin(x), we choose integers a, b so that
    #
    #    x = s + a * (pi/2) + b * (pi/16)
    #
    #  and |s| <= pi/32.  Using the fact that
    #
    #    sin(pi/16) = 0.5 * sqrt(2 - sqrt(2 + sqrt(2)))
    #
    #  we can compute sin(x) from sin(s), cos(s).  This greatly
    #  increases the convergence of the sine Taylor series.

    if iszero(a)
        return zero(a)
    end

    # approximately reduce modulo 2*pi
    z = round(a / double_2pi)
    r = a - double_2pi * z

    # approximately reduce modulo pi/2 and then modulo pi/16.

    q = floor(r.hi / double_pi2.hi + 0.5)
    t = r - double_pi2 * q
    j = convert(Int, q)
    q = floor(t.hi / double_pi16.hi + 0.5)
    t -= double_pi16 * q
    k = convert(Int, q)
    abs_k = abs(k)

    if j < -2 || j > 2
        # Cannot reduce modulo pi/2.
        return double_nan
    end

    if (abs_k > 4)
        # Cannot reduce modulo pi/16.
        return double_nan
    end

    if k == 0
        if j == 0
            return sin_taylor(t)
        elseif j == 1
            return cos_taylor(t)
        elseif j == -1
            return -cos_taylor(t)
        else
            return -sin_taylor(t)
        end
    end

    u = cos_table[abs_k]
    v = sin_table[abs_k]
    sin_t, cos_t = sincos_taylor(t)

    if j == 0
        r = k > 0 ? u * sin_t + v * cos_t : u * sin_t - v * cos_t
    elseif j == 1
        r = k > 0 ? u * cos_t - v * sin_t : u * cos_t + v * sin_t
    elseif j == -1
        r = k > 0 ? v * sin_t - u * cos_t : -u * cos_t - v * sin_t
    else
        r = k > 0 ? -u * sin_t - v * cos_t : v * cos_t - u * sin_t
    end

    r
end

function Base.cos(a::DoubleF64)
    if iszero(a)
        return one(a)
    end

    # approximately reduce modulo 2*pi
    z = round(a / double_2pi)
    r = a - double_2pi * z

    # approximately reduce modulo pi/2 and then modulo pi/16.

    q = floor(r.hi / double_pi2.hi + 0.5)
    t = r - double_pi2 * q
    j = convert(Int, q)
    q = floor(t.hi / double_pi16.hi + 0.5)
    t -= double_pi16 * q
    k = convert(Int, q)
    abs_k = abs(k)

    if j < -2 || j > 2
        # Cannot reduce modulo pi/2.
        return double_nan
    end

    if (abs_k > 4)
        # Cannot reduce modulo pi/16.
        return double_nan
    end

    if k == 0
        if j == 0
            return cos_taylor(t)
        elseif j == 1
            return -sin_taylor(t)
        elseif j == -1
            return sin_taylor(t)
        else
            return -cos_taylor(t)
        end
    end

    u = cos_table[abs_k]
    v = sin_table[abs_k]
    sin_t, cos_t = sincos_taylor(t)

    if j == 0
        r = k > 0 ? u * cos_t - v * sin_t : u * cos_t + v * sin_t
    elseif j == 1
        r = k > 0 ? -u * sin_t - v * cos_t : v * cos_t - u * sin_t
    elseif j == -1
        r = k > 0 ? u * sin_t + v * cos_t : u * sin_t - v * cos_t
    else
        r = k > 0 ? v * sin_t - u * cos_t : -u * cos_t - v * sin_t
    end

    r
end

function Base.sincos(a::DoubleF64)
    if iszero(a)
        return zero(a), one(a)
    end

    # approximately reduce modulo 2*pi
    z = round(a / double_2pi)
    r = a - double_2pi * z

    # approximately reduce modulo pi/2 and then modulo pi/16.
    q = floor(r.hi / double_pi2.hi + 0.5)
    t = r - double_pi2 * q
    j = convert(Int, q)
    abs_j = abs(j)
    q = floor(t.hi / double_pi16.hi + 0.5)
    t -= double_pi16 * q
    k = convert(Int, q)
    abs_k = abs(k)

    if abs_j > 2
        # Cannot reduce modulo pi/2.
        return double_nan, double_nan
    end

    if abs_k > 4
        # Cannot reduce modulo pi/16.
        return double_nan, double_nan
    end

    sin_t, cos_t = sincos_taylor(t)

    if abs_k == 0
        s = sin_t
        c = cos_t
    else
        u = cos_table[abs_k]
        v = sin_table[abs_k]

        s = k > 0 ? u * sin_t + v * cos_t : u * sin_t - v * cos_t
        c = k > 0 ? u * cos_t - v * sin_t : u * cos_t + v * sin_t
    end

    if j == 0
        s, c
    elseif j == 1
        c, -s
    elseif j == -1
        -c, s
    else
        -s, -c
    end
end

Base.atan(a::DoubleF64) = atan(a, one(a))

function Base.atan(y::DoubleF64, x::DoubleF64)
    # Strategy: Instead of using Taylor series to compute
    # arctan, we instead use Newton's iteration to solve
    # the equation
    #
    #    sin(z) = y/r    or    cos(z) = x/r
    #
    # where r = sqrt(x^2 + y^2).
    # The iteration is given by
    #
    #    z' = z + (y - sin(z)) / cos(z)          (for equation 1)
    #    z' = z - (x - cos(z)) / sin(z)          (for equation 2)
    #
    # Here, x and y are normalized so that x^2 + y^2 = 1.
    # If |x| > |y|, then first iteration is used since the
    # denominator is larger.  Otherwise, the second is used.
    if iszero(x)
        if iszero(y)
            double_nan
        end

        return y.hi > 0.0 ? double_pi2 : -double_pi2
    elseif iszero(y)
        return x.hi > 0.0 ? zero(x) : double_pi
    end

    if x == y
        return y.hi > 0.0 ? double_pi4 : -double_3pi4
    end

    if x == -y
        return y.hi > 0.0 ? double_3pi4 : -double_pi4
    end

    r = sqrt(square(x) + square(y))
    xx = x / r
    yy = y / r

    # Compute double precision approximation to atan.
    z = DoubleF64(atan(convert(Float64, y), convert(Float64, x)))

    if abs(xx.hi) > abs(yy.hi)
        # Use Newton iteration 1.  z' = z + (y - sin(z)) / cos(z)
        sin_z, cos_z = sincos(z)
        z += (yy - sin_z) / cos_z
    else
        # Use Newton iteration 2.  z' = z - (x - cos(z)) / sin(z)
        sin_z, cos_z = sincos(z)
        z -= (xx - cos_z) / sin_z
    end

    z
end

function Base.tan(a::DoubleF64)
    s, c = sincos(a)
    s / c
end

function Base.asin(a::DoubleF64)
    abs_a = abs(a)

    if abs_a > 1.0
        throw(DomainError())
    end

    if isone(abs_a)
        return a.hi > 0.0 ? double_pi2 : -double_pi2
    end

    atan(a, sqrt(1.0 - square(a)))
end

function Base.acos(a::DoubleF64)
    abs_a = abs(a)

    if abs_a > 1.0
        throw(DomainError())
    end

    if isone(abs_a)
        return a.hi > 0.0 ? zero(a) : -double_pi
    end

    atan(sqrt(1.0 - square(a)), a)
end

function Base.show(io::IO, x::DoubleF64)
    Printf.@printf io "%.32g" big(x)  # crude approximation to valid number of digits
end

end # module


using Base.Test, Whittaker

@testset "whittaker1" begin
  @testset "out-of-place" begin
    tol = 0.0001

    y = [1.0, 1.2, 0, 1.6, 1.8]
    w = [1.0, 1.0, 0, 1.0, 1.0]
    lambda = 0.0001
    z = whittaker1(y, w, lambda)
    @test_approx_eq_eps z [1.00002, 1.2, 1.4, 1.6, 1.79998] tol

    y = [1., 2, 0, 4, 5]
    w = [1, 1, 0, 1, 1]
    lambda = 0.0001
    z = whittaker1(y, w, lambda)
    @test_approx_eq_eps z [1.0001, 2.0, 3.0, 4.0, 4.9999] tol
  end

  @testset "in-place" begin
    tol = 0.0001

    y = [1.0, 1.2, 0, 1.6, 1.8]
    w = [1.0, 1.0, 0, 1.0, 1.0]
    lambda = 0.0001
    whittaker1!(y, w, lambda)
    @test_approx_eq_eps y [1.00002, 1.2, 1.4, 1.6, 1.79998] tol

    y = [1., 2, 0, 4, 5]
    w = [1, 1, 0, 1, 1]
    lambda = 0.0001
    whittaker1!(y, w, lambda)
    @test_approx_eq_eps y [1.0001, 2.0, 3.0, 4.0, 4.9999] tol
  end
end


"""
    Hyperbolic{T <: Number} <: Construct{T}

A hyperbolic Cayley-Dickson construct as a pair.
"""
struct Hyperbolic{T <: Number} <: Construct{T}
    l::T
    r::T
end

function (==)(x::Hyperbolic, y::Hyperbolic)
    x.l == y.l && x.r == y.r
end

function real(z::Hyperbolic)
    real(z.l)
end

"""
    quadrance(z::CayleyDickson.Hyperbolic)

For hyperbolic Cayley-Dickson constructs, return the difference of the squared
absolute value of the left and right components of `z`.
Note that this can be positive, negative, or zero.
"""
function quadrance(z::Hyperbolic)
    quadrance(z.l) - quadrance(z.r)
end

function iszerodivisor(z::Hyperbolic)
    quadrance(z.l) == quadrance(z.r)
end

function zero(::Type{Hyperbolic{T}}) where T <: Number
    Hyperbolic{T}(zero(T), zero(T))
end

function zero(::Hyperbolic{T}) where T <: Number
    zero(Hyperbolic{T})
end

function one(::Type{Hyperbolic{T}}) where T <: Number
    Hyperbolic{T}(one(T), zero(T))
end

function one(::Hyperbolic{T}) where T <: Number
    one(Hyperbolic{T})
end

function (-)(z::Hyperbolic)
    Hyperbolic(-(z.l), -(z.r))
end

function conj(z::Hyperbolic)
    Hyperbolic(conj(z.l), -(z.r))
end

function (+)(x::Hyperbolic{T}, y::Hyperbolic{T}) where T <: Number
    Hyperbolic{T}(x.l + y.l, x.r + y.r)
end

function (+)(a::T, z::Hyperbolic{T}) where T <: Number
    Hyperbolic{T}(a + z.l, z.r)
end

function (+)(z::Hyperbolic{T}, a::T) where T <: Number
    Hyperbolic{T}(z.l + a, z.r)
end

function (+)(a::Real, z::Hyperbolic)
    l = a + z.l
    Hyperbolic(l, oftype(l, z.r))
end

function (+)(z::Hyperbolic, a::Real)
    l = z.l + a
    Hyperbolic(l, oftype(l, z.r))
end

function (-)(x::Hyperbolic{T}, y::Hyperbolic{T}) where T <: Number
    Hyperbolic{T}(x.l - y.l, x.r - y.r)
end

function (-)(a::T, z::Hyperbolic{T}) where T <: Number
    Hyperbolic{T}(a - z.l, -(z.r))
end

function (-)(z::Hyperbolic{T}, a::T) where T <: Number
    Hyperbolic{T}(z.l - a, z.r)
end

function (-)(a::Real, z::Hyperbolic)
    l = a - z.l
    Hyperbolic(l, -(oftype(l, z.r)))
end

function (-)(z::Hyperbolic, a::Real)
    l = z.l - a
    Hyperbolic(l, oftype(l, z.r))
end

function (*)(x::Hyperbolic{T}, y::Hyperbolic{T}) where T <: Number
    Hyperbolic{T}(
        (x.l * y.l) + (conj(y.r) * x.r),
        (y.r * x.l) + (x.r * conj(y.l))
    )
end

function (*)(a::T, z::Hyperbolic{T}) where T <: Number
    Hyperbolic{T}(a * z.l, z.r * a)
end

function (*)(z::Hyperbolic{T}, a::T) where T <: Number
    Hyperbolic{T}(z.l * a, z.r * conj(a))
end

function (*)(a::Real, z::Hyperbolic)
    Hyperbolic(a * z.l, a * z.r)
end

function (*)(z::Hyperbolic, a::Real)
    Hyperbolic(z.l * a, z.r * a)
end

function (/)(z::Hyperbolic, a::Real)
    if iszero(a)
        error(ZeroDenominator)
    end

    Hyperbolic(z.l / a, z.r / a)
end

function (\)(a::Real, z::Hyperbolic)
    if iszero(a)
        error(ZeroDenominator)
    end

    Hyperbolic(a \ z.l, a \ z.r)
end

function widen(::Type{Hyperbolic{T}}) where T
    Hyperbolic{widen(T)}
end

function float(::Type{Hyperbolic{T}}) where T <: AbstractFloat
    Hyperbolic{T}
end

function float(::Type{Hyperbolic{T}}) where T
    Hyperbolic{float(T)}
end

"""
    cloak(a::Hyperbolic)

For hyperbolic Cayley-Dickson constructs, equivalent to the minus operation.
"""
function cloak(a::Hyperbolic)
    -(a)
end

"""
    dagger(a::Hyperbolic)

For hyperbolic Cayley-Dickson constructs, equivalent to the identity operation.
"""
function dagger(a::Hyperbolic)
    a
end

"""
    star(a::Hyperbolic)

For hyperbolic Cayley-Dickson constructs, the identity operation.
"""
function star(a::Hyperbolic)
    a
end


@add_problem constraints_benchmark function LT_constraints(handle_problem!, args...)
    p = satisfy()
    x = [Variable() for _ = 1:1000]
    for (i, xi) in enumerate(x)
        push!(p.constraints, xi <= 1.0 * i)
    end
    handle_problem!(p)
    nothing
end

@add_problem constraints_benchmark function LT_constraint(handle_problem!, args...)
    p = satisfy()
    x = Variable(1000)
    push!(p.constraints, x <= collect(1.0:1000.0))
    handle_problem!(p)
    nothing
end

@add_problem constraints_benchmark function GT_constraints(handle_problem!, args...)
    p = satisfy()
    x = [Variable() for _ = 1:1000]
    for (i, xi) in enumerate(x)
        push!(p.constraints, xi >= 1.0 * i)
    end
    handle_problem!(p)
    nothing
end

@add_problem constraints_benchmark function GT_constraint(handle_problem!, args...)
    p = satisfy()
    x = Variable(1000)
    push!(p.constraints, x >= collect(1.0:1000.0))
    handle_problem!(p)
    nothing
end

@add_problem constraints_benchmark function equality_constraints(handle_problem!, args...)
    p = satisfy()
    x = [Variable() for _ = 1:1000]
    for (i, xi) in enumerate(x)
        push!(p.constraints, xi == 1.0 * i)
    end
    handle_problem!(p)
    nothing
end

@add_problem constraints_benchmark function equality_constraint(handle_problem!, args...)
    p = satisfy()
    x = Variable(1000)
    push!(p.constraints, x == collect(1.0:1000.0))
    handle_problem!(p)
    nothing
end

@add_problem constraints_benchmark function sdp_constraint(handle_problem!, args...)
    p = satisfy()
    x = Variable(44, 44) # 990 vectorized entries
    push!(p.constraints, x ⪰ 0)
    handle_problem!(p)
    nothing
end

@add_problem constraints_benchmark function sdp_constraints(handle_problem!, args...)
    p = satisfy()
    x = [Variable(4, 4) for _ = 1:100] # 1000 total vectorized entries
    for v in x
        push!(p.constraints, v ⪰ 0)
    end
    handle_problem!(p)
    nothing
end


function get_output(coefs,Df,sq,r2,f_terms,Gres,Tot,p,n)
    n_term = length(coefs)
    DF = RegCol("Df")
    R2 = RegCol("R²")
    sumsq = RegCol("SumOfSqs")
    F = RegCol("F")
    P = RegCol("P")

    for i in 1:n_term
        setcoef!(R2, coefs[i]=>r2[i])
        setcoef!(DF, coefs[i]=>Df[i])
        setcoef!(sumsq, coefs[i]=>sq[i])
        setcoef!(F, coefs[i]=>f_terms[i])
        setcoef!(P, coefs[i]=>p[i])
    end
    coefArray = hcat(Df,sq,r2,f_terms,p)
    coefArray= vcat(coefArray,[n-1-sum(Df) Gres 1-sum(r2) NaN NaN ])
    coefArray= vcat(coefArray,[n-1 Tot 1 NaN NaN ])

    setcoef!(R2, "Residual"=>1-sum(r2))
    setcoef!(DF, "Residual"=>n-1-sum(Df))
    setcoef!(sumsq, "Residual"=>Gres)
    setcoef!(R2, "Total"=>1)
    setcoef!(DF, "Total"=>n-1)
    setcoef!(sumsq, "Total"=>Tot)
    return NamedArray(coefArray,(vcat(coefs...,"Residual","Total"),["Df","SumOfSqs","R2","F","P"])), hcat(DF,sumsq,R2,F,P)
end


export batchnorm, bnmoments, bnparams
using AutoGrad: AutoGrad, @primitive1, value, recording
using Statistics: mean, var

# dy20200804: why not have bnparams and bnmoments in a single struct?
# dy20200804: what array types / sizes are they supposed to use? which ones are Param?
# moments and params need to have the same atype as x, params needs to be wrapped in Param.
#TODO: improve the documentation. Explain what batchnorm is.

mutable struct BNMoments
    momentum::AbstractFloat
    mean
    var
    meaninit
    varinit
end

"""
    batchnorm(x[, moments, params]; kwargs...)

perform batch normalization on `x` with optional mean and variance in `moments` and scaling
factor and bias in `params`. See https://arxiv.org/abs/1502.03167 for reference.

2d, 4d and 5d inputs are supported. Mean and variance are computed over dimensions (2,),
(1,2,4) and (1,2,3,5) for 2d, 4d and 5d arrays, respectively.

`moments` stores running mean and variance to be used at inference time.  It is optional in
training mode, but mandatory in test mode.  Training and test modes can be controlled by the
`training` keyword argument which defaults to `Knet.training()`.

`params` stores the optional affine parameters gamma and beta.  `bnparams` function can be
used to initialize `params`.

# Example

    # Inilization, C is an integer
    moments = bnmoments()
    params = bnparams(C)
    ...
    # size(x) -> (H, W, C, N)
    y = batchnorm(x, moments, params)
    # size(y) -> (H, W, C, N)

# Keywords

 `eps=1e-5`: The epsilon parameter added to the variance to avoid division by 0.

 `training=Knet.training()`: When `training` is true, the mean and variance of `x` are used
 and `moments` argument is modified if it is provided. When `training` is false, mean and
 variance stored in the `moments` argument are used.

"""
function batchnorm(x, moments::Union{BNMoments, Nothing}=nothing, params=nothing;
                   training=AutoGrad.recording(), o...)
    xnd = ndims(x)
    a = (x,)
    if params !== nothing
        g = reshape(_bnscale(params), _wsize(x))
        b = reshape(_bnbias(params), _wsize(x))
        a = (g, b, x)
    end
    if xnd == 2
        return batchnorm2(a...; o...,
                          moments=moments,
                          training=training,
                          cache=BNCache())
    elseif xnd in [4, 5]
        return batchnorm4(a...; o...,
                          moments=moments,
                          training=training,
                          cache=BNCache())
    else
        error("Unsupported input dimension ", xnd)
    end
end

"""
    bnmoments(;momentum=0.1, mean=nothing, var=nothing, meaninit=zeros, varinit=ones)

Return a `BNMoments` object, a data structure used to store running mean and running
variance of batch normalization with the following fields:

* `momentum=0.1`: A real number between 0 and 1 to be used as the scale of last
mean and variance. The existing running mean or variance is multiplied by (1-momentum).

* `mean=nothing`: The running mean.

* `var=nothing`: The running variance.

* `meaninit=zeros`: The function used for initialize the running mean. Should either be
`nothing` or of the form `([eltype], dims...)->data`. `zeros` is a good option.

* `varinit=ones`: The function used for initialize the running variance. Should either be
`nothing` or `([eltype], dims...)->data`. `ones` is a good option.

This constructor can be used directly load moments from data. `meaninit` and `varinit` are
called if `mean` and `var` are nothing. Type and size of the `mean` and `var` are determined
automatically from the inputs in the `batchnorm` calls.

"""
bnmoments(;momentum=0.1, meaninit=zeros, varinit=ones, mean=nothing, var=nothing) =
    BNMoments(momentum, mean, var, meaninit, varinit)

"""
    bnparams(etype, channels::Integer)

Return a single 1d array that contains both scale and bias of batchnorm, where the first
half is scale and the second half is bias.

`bnparams(channels)` calls `bnparams(Float64, channels)`, following Julia convention.

"""
function bnparams(etype, channels::Integer)
    buf = Array{etype}(undef,2channels)
    buf[1:channels] .= 1
    buf[channels+1:end] .= 0
    return buf
end

bnparams(channels::Integer) = bnparams(Float64, channels)

#=
LOW-LEVEL API that won't be exported by default
=#
@inline _bnscale(param) = param[1:div(length(param), 2)]
@inline _bnbias(param) = param[div(length(param), 2)+1:end]

# Dimension helpers
@inline _wsize(y) = ((1 for _=1:ndims(y)-2)..., size(y)[end-1], 1)
@inline _reddims(y) = ((i for i=1:ndims(y)-2)..., ndims(y))

# A black box data type for storing the bn state
mutable struct BNCache
    mean
    ivar
    dx
    dg
    db
end

BNCache() = BNCache(nothing, nothing, nothing, nothing, nothing)

# CPU Implementation
function batchnorm4(g, b, x; o...)
    return _batchnorm4_fused(g, b, x; o...)
end

function batchnorm4(x; o...)
    return _batchnorm4_fused(nothing,nothing,x; o...)
end

function _batchnorm4_fused(g, b, x; eps=1e-5, training=AutoGrad.recording(), cache=nothing, moments=nothing, o...)
    T = eltype(x)
    y = copy(x)
    eps = T(eps)
    dims = _reddims(y)
    if moments!==nothing
        _lazy_init!(moments, x)
    end
    affine = (g !== nothing)
    if training
        mu = mean(y, dims=dims)
        sigma2 = var(y, dims=dims, corrected=false, mean=mu)
        _update_moments!(moments, mu, sigma2)
        sigma2 .= 1 ./ sqrt.(sigma2 .+ eps)
        ivar = sigma2
        if affine
            y .= g .* (y .- mu) .* ivar .+ b
        else
            y .= (y .- mu) .* ivar
        end
        if cache !== nothing
            cache.ivar = ivar
            cache.mean = mu
        end
    else
        @assert (moments!==nothing)  "Moments must be provided in test mode"
        ivar = 1 ./ sqrt.(moments.var .+ eps)
        if affine
            y .= g .* (y .- moments.mean) .* ivar .+ b
        else
            y .= (y .- moments.mean) .* ivar
        end
    end
    return y
end

function _update_moments!(moments, mu, sigma2)
    moments === nothing && return
    moments.mean .= moments.momentum .* mu .+
        (1 .- moments.momentum) .* moments.mean
    moments.var .= moments.momentum .* sigma2 .+
        (1 .- moments.momentum) .* moments.var
end

# TODO: consider automatic type conversion
function _lazy_init!(m::BNMoments, x)
    x = value(x)
    buf_size = (ndims(x) > 2) ? _wsize(x) : (size(x,1), 1)
    tx = typeof(x)
    ex = eltype(x)
    m.momentum = ex(m.momentum)
    if m.mean == nothing
        m.mean = tx(m.meaninit(ex, buf_size...))
    end
    if m.var == nothing
        m.var = tx(m.varinit(ex, buf_size...))
    end
end

# Implement batchnorm2 using batchnorm4, with autograd

function batchnorm2(g, b, x; moments=nothing, training=AutoGrad.recording(), o...)
    # TODO: This support should be added when needed
    if training == false && (isa(g, AutoGrad.Value) || isa(x, AutoGrad.Value) || isa(b, AutoGrad.Value))
        error("Test mode backward is not supported with 2d inputs")
    end
    @inline _pad4(x) = reshape(x, (1,1,size(x,1),size(x,2)))
     # process moments
    if moments !== nothing
        _lazy_init!(moments, x)
        moments.mean = _pad4(moments.mean)
        moments.var = _pad4(moments.var)
    end
    x = _pad4(x)
    args = (x,)
    if g !== nothing
        g = _pad4(g)
        b = _pad4(b)
        args = (g, b, x)
    end
    y = mat(batchnorm4(args...;
                       moments=moments,
                       training=training,
                       o...))
    if moments !== nothing
        moments.mean = mat(moments.mean)
        moments.var = mat(moments.var)
    end
    return y
end

batchnorm2(x;o...) = batchnorm2(nothing, nothing, x; o...)

# CPU backward
function batchnorm4_back(g, x, dy; eps=1e-5, training=AutoGrad.recording(), cache=nothing, moments=nothing,  o...)
    T = eltype(x)
    eps = T(eps)
    dims = _reddims(x)
    if training
        mu, ivar = _get_cache_data(cache, x, eps)
        x_mu = x .- mu
        # equations from the original paper
        dyivar = dy .* ivar
        if g !== nothing
            dg = sum(x_mu .* dyivar, dims=dims)
            db = sum(dy, dims=dims)
            dyivar .*= g
        else
            dg, db = nothing, nothing
        end
        m = prod(d->size(x,d), dims) # size(x, dims...))
        dsigma2 = -sum(dyivar .* x_mu .* ivar.^2, dims=dims) ./ 2
        dmu = -sum(dyivar, dims=dims) .- 2dsigma2 .* sum(x_mu, dims=dims) ./ m
        dx = dyivar .+ dsigma2 .* 2x_mu ./ m .+ dmu ./ m
    else #same reasoning with the gpu version
        ivar = 1 ./ sqrt.(moments.var .+ eps)
        dx = (g !== nothing) ? (dy .* g .* ivar) : (dy .* ivar)
        if g !== nothing
            dg = sum(dy .* (x .- moments.mean) .* ivar, dims=dims)
            db = sum(dy, dims=dims)
        else
            dg, db = nothing, nothing
        end
    end
    return dg, db, dx
end

function _get_cache_data(cache, x, eps)
    if cache !== nothing
        mu = cache.mean
        ivar = cache.ivar
    else
        mu = mean(x, dims=_reddims(x))
        ivar = 1 ./ sqrt.(var(x, dims=_reddims(x); mean=mu, correct=false) .+ eps)
    end
    return mu, ivar
end

function batchnorm4g(g,x,dy; cache=nothing, o...)
    dg, db, dx = batchnorm4_back(g, x, dy; cache=cache, o...)
    if cache !== nothing
        cache.dx, cache.db = dx, db
    else
        @warn("Calling batchnorm with affine without cache is not recommended for performance")
    end
    return dg
end

function batchnorm4b(dy; cache=nothing, o...)
    if cache == nothing || cache.db == nothing
        return sum(dy, dims=_reddims(dy))
    else
        return cache.db
    end
end

function batchnorm4x(g, x, dy; cache=nothing, o...)
    if cache !== nothing && cache.dx !== nothing
        return cache.dx
    else
        return batchnorm4_back(g, x, dy; cache=cache, o...)[3]
    end
end

function batchnorm4x(x, dy; o...)
    return batchnorm4_back(nothing, x, dy; o...)[3]
end

@primitive1 batchnorm4(x;o...),dy batchnorm4x(x, dy; o...)

# 4d with affine
@primitive1 batchnorm4(g, b, x; o...),dy batchnorm4g(g, x, dy; o...) batchnorm4b(dy; o...) batchnorm4x(g, x, dy; o...)