import Base: ==, +, -, *, /, merge, merge!

#for StatsBase histogram
for op in (:+, :-, :*, :/)
    @eval ($op)(h1::Histogram, h2::Histogram) =
    (==)(h1.edges,h2.edges) ? Histogram(h1.edges, broadcast($op, h1.weights, h2.weights)) : throw(DimensionMismatch("The bins of the two histograms don't match"))
    @eval ($op)(h1::Histogram, n::Real) = Histogram(h1.edges, broadcast($op, h1.weights, n))
end

for T in (:Hist1D,:Hist2D)
    for op in (:+, :-)
        @eval function ($op)(h1::($T), h2::($T))
            h1.hist.edges != h2.hist.edges && throw(DimensionMismatch("Edges don't match"))
            _f(counts) = any(x -> x<0, counts)
            _hist = ($op)(h1.hist,  h2.hist)
            ($T)(_hist, h1.sumw2 + h2.sumw2)
        end
    end

    @eval function *(h1::($T), num::Real)
        _f(counts) = any(x -> x<0, counts)
        _f(h1.hist.weights) && error("Can't do * when bin count is negative")
        _hist = *(h1.hist, num)
        ($T)(_hist, h1.sumw2 * num^2)
    end

    # https://github.com/aminnj/yahist/blob/4a5767f181ec7fdcc4af18cf15ceedd1c2f89019/yahist/hist1d.py#L427-L430
    @eval function /(h1::($T), h2::($T))
        _f(counts) = any(x -> x<0, counts)
        (_f(h1.hist.weights) || _f(h2.hist.weights)) && error("Can't do / when bin count is negative")
        _hist = /(h1.hist, h2.hist)
        _sumw2 =  @. h1.sumw2 / (h2.hist.weights ^ 2) +
                (sqrt(h2.sumw2) * h1.hist.weights / (h2.hist.weights ^ 2)) ^ 2

        ($T)(_hist, _sumw2)
    end

    @eval function ==(h1::$T, h2::$T)
        h1.hist == h2.hist &&
        h1.sumw2 == h2.sumw2
    end

    @eval function merge!(h1::$T, h2::$T)
        h1.hist.edges != h2.hist.edges && throw(DimensionMismatch("The dimension doesn't match"))
        lock(h1)
        h1.hist.weights .+= h2.hist.weights
        h1.sumw2 .+= h2.sumw2
        unlock(h1)
        ($T)(h1.hist)
    end

    @eval merge(h1::$T, h2::$T) = merge!(deepcopy(h1), h2)
end

function merge(hists...)
    h = deepcopy(first(hists))
    length(hists) == 1 && return h
    foreach(x-> merge!(h, x), hists[2:end])
    h
end


export IntersectionArray,
       array

# ================================
# Intersection of an array of sets
# ================================

"""
    IntersectionArray{N, S<:LazySet{N}} <: LazySet{N}

Type that represents the intersection of a finite number of sets.

### Fields

- `array` -- array of sets

### Notes

This type assumes that the dimensions of all elements match.

The `EmptySet` is the absorbing element for `IntersectionArray`.

The intersection preserves convexity: if the set arguments are convex, then
their intersection is convex as well.

Constructors:

- `IntersectionArray(array::Vector{<:LazySet})` -- default constructor

- `IntersectionArray([n]::Int=0, [N]::Type=Float64)`
 -- constructor for an empty sum with optional size hint and numeric type
"""
struct IntersectionArray{N, S<:LazySet{N}} <: LazySet{N}
   array::Vector{S}
end

isoperationtype(::Type{<:IntersectionArray}) = true
isconvextype(::Type{IntersectionArray{N, S}}) where {N, S} = isconvextype(S)

# constructor for an empty sum with optional size hint and numeric type
function IntersectionArray(n::Int=0, N::Type=Float64)
   arr = Vector{LazySet{N}}()
   sizehint!(arr, n)
   return IntersectionArray(arr)
end

# Universe is the neutral element for IntersectionArray
@neutral(IntersectionArray, Universe)

# EmptySet is the absorbing element for IntersectionArray
@absorbing(IntersectionArray, EmptySet)

# add functions connecting Intersection and IntersectionArray
@declare_array_version(Intersection, IntersectionArray)

"""
    array(ia::IntersectionArray)

Return the array of an intersection of a finite number of sets.

### Input

- `ia` -- intersection of a finite number of sets

### Output

The array of an intersection of a finite number of sets.
"""
function array(ia::IntersectionArray)
   return ia.array
end

# --- LazySet interface functions ---

"""
    dim(ia::IntersectionArray)

Return the dimension of an intersection of a finite number of sets.

### Input

- `ia` -- intersection of a finite number of sets

### Output

The ambient dimension of the intersection of a finite number of sets.
"""
function dim(ia::IntersectionArray)
   return length(ia.array) == 0 ? 0 : dim(ia.array[1])
end

"""
    σ(d::AbstractVector, ia::IntersectionArray)

Return the support vector of an intersection of a finite number of sets in a
given direction.

### Input

- `d`  -- direction
- `ia` -- intersection of a finite number of sets

### Output

The support vector in the given direction.
If the direction has norm zero, the result depends on the individual sets.
"""
function σ(d::AbstractVector, ia::IntersectionArray)
    X = concretize(ia)
    return σ(d, X)
end

"""
    isbounded(ia::IntersectionArray)

Determine whether an intersection of a finite number of sets is bounded.

### Input

- `ia` -- intersection of a finite number of sets

### Output

`true` iff the intersection is bounded.

### Algorithm

We first check if any of the wrapped sets is bounded.
Otherwise, we check boundedness via [`LazySets._isbounded_unit_dimensions`](@ref).
"""
function isbounded(ia::IntersectionArray)
   if any(isbounded, ia.array)
       return true
   end
   return _isbounded_unit_dimensions(ia)
end

function isboundedtype(::Type{<:IntersectionArray{N, S}}) where {N, S}
    return isboundedtype(S)
end

"""
    ∈(x::AbstractVector, ia::IntersectionArray)

Check whether a given point is contained in an intersection of a finite number
of sets.

### Input

- `x`  -- point/vector
- `ia` -- intersection of a finite number of sets

### Output

`true` iff ``x ∈ ia``.
"""
function ∈(x::AbstractVector, ia::IntersectionArray)
   for S in ia.array
       if x ∉ S
           return false
       end
   end
   return true
end

"""
    constraints_list(ia::IntersectionArray{N}) where {N}

Return the list of constraints of an intersection of a finite number of
(polyhedral) sets.

### Input

- `ia` -- intersection of a finite number of (polyhedral) sets

### Output

The list of constraints of the intersection.

### Notes

We assume that the underlying sets are polyhedral, i.e., offer a method
`constraints_list`.

### Algorithm

We create the polyhedron from the `constraints_list`s of the sets and remove
redundant constraints.
"""
function constraints_list(ia::IntersectionArray{N}) where {N}
   constraints = Vector{LinearConstraint{N, Vector{N}}}() # TODO: use vector type of ia
   for X in array(ia)
       clist_X = _normal_Vector(X)
       append!(constraints, clist_X)
   end
   remove_redundant_constraints!(constraints)
   return constraints
end

function concretize(ia::IntersectionArray)
    a = array(ia)
    @assert !isempty(a) "an empty intersection is not allowed"
    X = ia
    @inbounds for (i, Y) in enumerate(a)
        if i == 1
            X = concretize(Y)
        else
            X = intersection(X, concretize(Y))
        end
    end
    return X
end


# # [A regression task: the prediction of  bike  sharing demand](@id regression_tutorial)
# The task is to estimate the influence of several variables (like the weather, the season, the day of the week..) on the demand of shared bicycles, so that the authority in charge of the service can organise the service in the best way.
#
# Data origin:
# - original full dataset (by hour, not used here): [https://archive.ics.uci.edu/ml/datasets/Bike+Sharing+Dataset](https://archive.ics.uci.edu/ml/datasets/Bike+Sharing+Dataset)
# - simplified dataset (by day, with some simple scaling): [https://www.hds.utc.fr/~tdenoeux/dokuwiki/en/aec]( https://www.hds.utc.fr/~tdenoeux/dokuwiki/en/aec)
# - description: [https://www.hds.utc.fr/~tdenoeux/dokuwiki/_media/en/exam_2019_ace_.pdf](https://www.hds.utc.fr/~tdenoeux/dokuwiki/_media/en/exam_2019_ace_.pdf)
# - data: [https://www.hds.utc.fr/~tdenoeux/dokuwiki/_media/en/bike_sharing_day.csv.zip](https://www.hds.utc.fr/~tdenoeux/dokuwiki/_media/en/bike_sharing_day.csv.zip)

# Note that even if we are estimating a time serie, we are not using here a recurrent neural network as we assume the temporal dependence to be negligible (i.e. $Y_t = f(X_t)$ alone).

# ## Library and data loading

# We first load all the packages we are going to use
using  LinearAlgebra, Random, Statistics, DataFrames, CSV, Plots, Pipe, BenchmarkTools, BetaML
import Distributions: Uniform
import DecisionTree, Flux ## For comparisions
using  Test     #src

# Here we load the data from a csv provided by the BataML package
baseDir = joinpath(dirname(pathof(BetaML)),"..","docs","src","tutorials","Regression - bike sharing")
data    = CSV.File(joinpath(baseDir,"data","bike_sharing_day.csv"),delim=',') |> DataFrame
describe(data)

# The variable we want to learn to predict is `cnt`, the total demand of bikes for a given day. Even if it is indeed an integer, we treat it as a continuous variable, so each single prediction will be a scalar $Y \in \mathbb{R}$.
plot(data.cnt, title="Daily bike sharing rents (2Y)", label=nothing)

# ## Decision Trees

# We start our regression task with Decision Trees.

# Decision trees training consist in choosing the set of questions (in a hierarcical way, so to form indeed a "decision tree") that "best" split the dataset given for training, in the sense that the split generate the sub-samples (always 2 subsamples in the BetaML implementation) that are, for the characteristic we want to predict, the most homogeneous possible. Decision trees are one of the few ML algorithms that has an intuitive interpretation and can be used for both regression or classification tasks.

# ### Data preparation
# The first step is to prepare the data for the analysis. This indeed depends already on the model we want to employ, as some models "accept" almost everything as input, no matter if the data is numerical or categorical, if it has missing values or not... while other models are instead much more exigents, and require more work to "clean up" our dataset.

# Here we start using  Decision Tree and Random Forest models that definitly belong to the first group, so the only thing we have to do is to select the variables in input (the "feature matrix", that we will indicate with "X") and the variable representing our output (the information we want to learn to predict, we call it "y"):
x    = Matrix{Float64}(data[:,[:instant,:season,:yr,:mnth,:holiday,:weekday,:workingday,:weathersit,:temp,:atemp,:hum,:windspeed]])
y    = data[:,16];

# We can now split the dataset between the data we will use for training the algorithm (`xtrain`/`ytrain`), those for selecting the hyperparameters (`xval`/`yval`) and finally those for testing the quality of the algoritm with the optimal hyperparameters (`xtest`/`ytest`). We use the `partition` function specifying the share we want to use for these three different subsets, here 75%, 12.5% and 12.5 respectively. As our data represents indeed a time serie, we want our model to be able to predict _future_ demand of bike sharing from _past_, observed rented bikes, so we do not shuffle the datasets as it would be the default.

((xtrain,xval,xtest),(ytrain,yval,ytest)) = partition([x,y],[0.75,0.125,1-0.75-0.125],shuffle=false)
(ntrain, nval, ntest) = size.([ytrain,yval,ytest],1)

# We can now "tune" our model so-called hyper-parameters, i.e. choose the best exogenous parameters of our algorithm, where "best" refers to some minimisation of a "loss" function between the true and the predicted values. We compute this loss function on a specific subset of data, that we call the "validation" subset (`xval` and `yval`).

# BetaML doesn't have a dedicated function for hyper-parameters optimisation, but it is easy to write some custom julia code, at least for a simple grid-based "search". Indeed one of the main reasons that a dedicated function exists in other Machine Learning libraries is that loops in other languages are slow, but this is not a problem in julia, so we can retain the flexibility to write the kind of hyper-parameter tuning that best fits our needs.

# Below is an example of a possible such function. Note there are more "elegant" ways to code it, but this one does the job. In particular, for simplicity, this hyper-paramerter tuning function just run multiple repetitions. In real world it is better to use [cross-validation](@ref crossValidation) in the hyper-parameter tuning, expecially when the observations are small. The [Clustering tutorial](@ref clustering_tutorial) shows an example on how to use `crossValidation`.

# We will see the various functions inside `tuneHyperParameters()` in a moment. For now let's going just to observe that `tuneHyperParameters` just loops over all the possible hyper-parameters and selects the ones where the error between `xval` and `yval` is minimised. For the meaning of the various hyper-parameter, consult the documentation of the [`buildTree`](@ref) and [`buildForest`](@ref) functions.
# The function uses multiple threads, so we calls `generateParallelRngs()` (in the `BetaML.Utils` submodule) to generate thread-safe random number generators and locks the comparision step.

function tuneHyperParameters(model,xtrain,ytrain,xval,yval;maxDepthRange=15:15,maxFeaturesRange=size(xtrain,2):size(xtrain,2),nTreesRange=20:20,βRange=0:0,minRecordsRange=2:2,repetitions=5,rng=Random.GLOBAL_RNG)
    ## We start with an infinitely high error
    bestRme         = +Inf
    bestMaxDepth    = 1
    bestMaxFeatures = 1
    bestMinRecords  = 2
    bestNTrees      = 1
    bestβ           = 0
    compLock        = ReentrantLock()

    ## Generate one random number generator per thread
    masterSeed = rand(rng,100:9999999999999) ## Some RNG have problems with very small seed. Also, the master seed has to be computed _before_ generateParallelRngs
    rngs = generateParallelRngs(rng,Threads.nthreads())

    ## We loop over all possible hyperparameter combinations...
    parLengths = (length(maxDepthRange),length(maxFeaturesRange),length(minRecordsRange),length(nTreesRange),length(βRange))
    Threads.@threads for ij in CartesianIndices(parLengths) ## This to avoid many nested for loops
           (maxDepth,maxFeatures,minRecords,nTrees,β)   = (maxDepthRange[Tuple(ij)[1]], maxFeaturesRange[Tuple(ij)[2]], minRecordsRange[Tuple(ij)[3]], nTreesRange[Tuple(ij)[4]], βRange[Tuple(ij)[5]]) ## The specific hyperparameters of this nested loop
           tsrng = rngs[Threads.threadid()] ## The random number generator is specific for each thread..
           joinedIndx = LinearIndices(parLengths)[ij]
           ## And here we make the seeding depending on the id of the loop, not the thread: hence we get the same results indipendently of the number of threads
           Random.seed!(tsrng,masterSeed+joinedIndx*10)
           totAttemptError = 0.0
           ## We run several repetitions with the same hyperparameter combination to account for stochasticity...
           for r in 1:repetitions
              if model == "DecisionTree"
                 ## Here we train the Decition Tree model
                 myTrainedModel = buildTree(xtrain,ytrain, maxDepth=maxDepth,maxFeatures=maxFeatures,minRecords=minRecords,rng=tsrng)
              else
                 ## Here we train the Random Forest model
                 myTrainedModel = buildForest(xtrain,ytrain,nTrees,maxDepth=maxDepth,maxFeatures=maxFeatures,minRecords=minRecords,β=β,rng=tsrng)
              end
              ## Here we make prediciton with this trained model and we compute its error
              ŷval   = predict(myTrainedModel, xval,rng=tsrng)
              rmeVal = meanRelError(ŷval,yval,normRec=false)
              totAttemptError += rmeVal
           end
           avgAttemptedDepthError = totAttemptError / repetitions
           begin
               lock(compLock) ## This step can't be run in parallel...
               try
                   ## Select this specific combination of hyperparameters if the error is the lowest
                   if avgAttemptedDepthError < bestRme
                     bestRme         = avgAttemptedDepthError
                     bestMaxDepth    = maxDepth
                     bestMaxFeatures = maxFeatures
                     bestNTrees      = nTrees
                     bestβ           = β
                     bestMinRecords  = minRecords
                   end
               finally
                   unlock(compLock)
               end
           end
    end
    return (bestRme,bestMaxDepth,bestMaxFeatures,bestMinRecords,bestNTrees,bestβ)
end

# We can now run the hyperparameter optimisation function with some "reasonable" ranges. To obtain replicable results we call `tuneHyperParameters` with `rng=copy(FIXEDRNG)`, where `FIXEDRNG` is a fixed-seeded random number generator guaranteed to maintain the same stream of random numbers even between different julia versions. That's also what we use for our unit tests (see the [`Getting started`](@ref dealing_with_stochasticity) for more details).
(bestRme,bestMaxDepth,bestMaxFeatures,bestMinRecords) = tuneHyperParameters("DecisionTree",xtrain,ytrain,xval,yval,
           maxDepthRange=3:7,maxFeaturesRange=10:12,minRecordsRange=2:6,repetitions=10,rng=copy(FIXEDRNG))
println("DT: $bestRme - $bestMaxDepth - $bestMaxFeatures - $bestMinRecords") #src

# Now that we have found the "optimal" hyperparameters we can build ("train") our model using them:
myTree = buildTree(xtrain,ytrain, maxDepth=bestMaxDepth, maxFeatures=bestMaxFeatures,minRecords=bestMinRecords,rng=copy(FIXEDRNG));

# Let's benchmark the time and memory usage of the training step of a decision tree:
@btime  buildTree(xtrain,ytrain, maxDepth=bestMaxDepth, maxFeatures=bestMaxFeatures,minRecords=bestMinRecords,rng=copy(FIXEDRNG));
# Individual decision trees are blazing fast, among the fastest algorithms we could use.

#-

# The above [`buildTree`](@ref) function produces a DecisionTree object that can be used to make predictions given some new features, i.e. given some X matrix of (number of observations x dimensions), predict the corresponding Y vector of scalers in R.
(ŷtrain,ŷval,ŷtest) = predict.([myTree], [xtrain,xval,xtest])

# Note that the above code uses the "dot syntax" to "broadcast" `predict()` over an array of label matrices. It is exactly equivalent to:
ŷtrain = predict(myTree, xtrain);
ŷval   = predict(myTree, xval);
ŷtest  = predict(myTree, xtest);

# We now compute the relative mean error for the training, the validation and the test set. The [`meanRelError`](@ref) is a very flexible error function. Without additional parameter, it computes, as the name says, the _mean relative error_, also known as the "mean absolute percentage error" ([MAPE](https://en.wikipedia.org/wiki/Mean_absolute_percentage_error)) between an estimated and a true vector.
# However it can also compute the _relative mean error_ (as we do here), or use a p-norm higher than 1.
# The _mean relative error_ enfatises the relativeness of the error, i.e. all observations and dimensions weigth the same, wether large or small. Conversly, in the _relative mean error_ the same relative error on larger observations (or dimensions) weights more.
# In this exercise we use the later, as our data has clearly some outlier days with very small rents, and we care more of avoiding our customers finding empty bike racks than having unrented bikes on the rack. Targeting a low mean average error would push all our predicitons down to try accomodate the low-level predicitons (to avoid a large relative error), and that's not what we want.

# For example let's consider the following example:
y     = [30,28,27,3,32,38];
ŷpref = [32,30,28,10,31,40];
ŷbad  = [29,25,24,5,28,35];

# Here ŷpref is an ipotetical output of a model that minimises the relative mean error, while ŷbad minimises the mean realative error.
meanRelError.([ŷbad, ŷpref],[y,y],normRec=true) ## Mean relative error
#-
meanRelError.([ŷbad, ŷpref],[y,y],normRec=false) ## Relative mean error
#-
plot([y ŷbad ŷpref], colour=[:black :red :green], label=["obs" "bad est" "good est"])

# We can then compute the relative mean error for the decision tree
(rmeTrain, rmeVal, rmeTest) = meanRelError.([ŷtrain,ŷval,ŷtest],[ytrain,yval,ytest],normRec=false)
#src (0.12659039939672598, 0.0937634812419113, 0.22233930540215602)
#-
@test rmeTest <= 0.24 #src

# We can plot the true labels vs the estimated one for the three subsets...
scatter(ytrain,ŷtrain,xlabel="daily rides",ylabel="est. daily rides",label=nothing,title="Est vs. obs in training period (DT)")
#-
scatter(yval,ŷval,xlabel="daily rides",ylabel="est. daily rides",label=nothing,title="Est vs. obs in validation period (DT)")
#-
scatter(ytest,ŷtest,xlabel="daily rides",ylabel="est. daily rides",label=nothing,title="Est vs. obs in testing period (DT)")

# Or we can visualise the true vs estimated bike shared on a temporal base.
# First on the full period (2 years) ...
ŷtrainfull = vcat(ŷtrain,fill(missing,nval+ntest))
ŷvalfull   = vcat(fill(missing,ntrain), ŷval, fill(missing,ntest))
ŷtestfull  = vcat(fill(missing,ntrain+nval), ŷtest)
plot(data[:,:dteday],[data[:,:cnt] ŷtrainfull ŷvalfull ŷtestfull], label=["obs" "train" "val" "test"], legend=:topleft, ylabel="daily rides", title="Daily bike sharing demand observed/estimated across the\n whole 2-years period (DT)")

# ..and then focusing on the testing period
stc = 620
endc = size(x,1)
plot(data[stc:endc,:dteday],[data[stc:endc,:cnt] ŷvalfull[stc:endc] ŷtestfull[stc:endc]], label=["obs" "val" "test"], legend=:bottomleft, ylabel="Daily rides", title="Focus on the testing period (DT)")

# The predictions aren't so bad in this case, however decision trees are highly instable, and the output could have depended just from the specific initial random seed.

# ## Random Forests
# Rather than trying to solve this problem using a single Decision Tree model, let's not try to use a _Random Forest_ model. Random forests average the results of many different decision trees and provide a more "stable" result.
# Being made of many decision trees, random forests are hovever more computationally expensive to train, but luckily they tend to self-tune (or self-regularise). In particular the parameters `maxDepth` and `maxFeatures` shouldn't need tuning.

# We still tune however the model for other parameters, and in particular the β parameter, a prerogative of BetaML Random Forests that allows to assign more weigth to the best performing trees in the forest. It may be particularly important if there are many outliers in the data we don't want to "learn" from.
minRecordsRange=[2,4,5]; nTreesRange=60:10:80; βRange=100:100:300
(bestRme,bestMaxDepth,bestMaxFeatures,bestMinRecords,bestNTrees,bestβ) = tuneHyperParameters("RandomForest",xtrain,ytrain,xval,yval,
        maxDepthRange=size(xtrain,1):size(xtrain,1),maxFeaturesRange=Int(round(sqrt(size(xtrain,2)))):Int(round(sqrt(size(xtrain,2)))),
        minRecordsRange=minRecordsRange,nTreesRange=nTreesRange,βRange=βRange,repetitions=5,rng=copy(FIXEDRNG))
println("RF: $bestRme $bestMinRecords $bestNTrees $bestβ") #src

# As for decision trees, once the hyper-parameters of the model are tuned we wan train again the model using the optimal parameters.
#src  bestNTrees=80; bestMinRecords=5; bestβ=100
#src bestNTrees=60; bestMinRecords=5; bestβ=200
myForest = buildForest(xtrain,ytrain, bestNTrees, maxDepth=bestMaxDepth,maxFeatures=bestMaxFeatures,minRecords=bestMinRecords,β=bestβ,oob=true,rng=copy(FIXEDRNG));

# Let's now benchmark the training of the BetaML Random Forest model
@btime buildForest(xtrain,ytrain, bestNTrees, maxDepth=bestMaxDepth,maxFeatures=bestMaxFeatures,minRecords=bestMinRecords,β=bestβ,oob=true,rng=copy(FIXEDRNG));
# Random forests are evidently slower than individual decision trees but are still relativly fast. We should also consider that they are by default efficiently parallelised, so their speed increases with the number of available cores (in building this documentation page, GitHub CI servers allow for a single core, so all the bechmark you see in this tutorial are run with a single core available).

#-

# Random forests support the so-called "out-of-bag" error, an estimation of the error that we would have when the model is applied on a testing sample.
# However in this case the oob reported is much smaller than the testing error we will actually find. This is due to the fact that the division between training/validation and testing in this exercise is not random, but has a temporal basis. It seems that in this example the data in validation/testing follows a different pattern/variance than those in training (in probabilistic terms, the daily observations are not i.i.d.).
oobError, trueTestMeanRelativeError  = myForest.oobError,meanRelError(ŷtest,ytest,normRec=true)
#+
(ŷtrain,ŷval,ŷtest)         = predict.([myForest], [xtrain,xval,xtest])
(rmeTrain, rmeVal, rmeTest) = meanRelError.([ŷtrain,ŷval,ŷtest],[ytrain,yval,ytest],normRec=false)
#src 0.06512088063750156, 0.09389025098838212, 0.22229683782214169
@test rmeTest <= 0.23 #src

# In this case we found an error very similar to the one employing a single decision tree. Let's print the observed data vs the estimated one using the random forest and then along the temporal axis:
scatter(ytrain,ŷtrain,xlabel="daily rides",ylabel="est. daily rides",label=nothing,title="Est vs. obs in training period (RF)")
#-
scatter(yval,ŷval,xlabel="daily rides",ylabel="est. daily rides",label=nothing,title="Est vs. obs in validation period (RF)")
#-
scatter(ytest,ŷtest,xlabel="daily rides",ylabel="est. daily rides",label=nothing,title="Est vs. obs in testing period (RF)")

# Full period plot (2 years):
ŷtrainfull = vcat(ŷtrain,fill(missing,nval+ntest))
ŷvalfull   = vcat(fill(missing,ntrain), ŷval, fill(missing,ntest))
ŷtestfull  = vcat(fill(missing,ntrain+nval), ŷtest)
plot(data[:,:dteday],[data[:,:cnt] ŷtrainfull ŷvalfull ŷtestfull], label=["obs" "train" "val" "test"], legend=:topleft, ylabel="daily rides", title="Daily bike sharing demand observed/estimated across the\n whole 2-years period (RF)")

# Focus on the testing period:
stc = 620
endc = size(x,1)
plot(data[stc:endc,:dteday],[data[stc:endc,:cnt] ŷvalfull[stc:endc] ŷtestfull[stc:endc]], label=["obs" "val" "test"], legend=:bottomleft, ylabel="Daily rides", title="Focus on the testing period (RF)")

# ### Comparison with DecisionTree.jl random forest
# We now compare our results with those obtained employing the same model in the [DecisionTree package](https://github.com/bensadeghi/DecisionTree.jl), using the default suggested hyperparameters:

# Hyperparameters of the DecisionTree.jl random forest model
n_subfeatures=-1; n_trees=bestNTrees; partial_sampling=1; max_depth=26
min_samples_leaf=bestMinRecords; min_samples_split=bestMinRecords; min_purity_increase=0.0; seed=3

# We train the model..
model = DecisionTree.build_forest(ytrain, convert(Matrix,xtrain),
                     n_subfeatures,
                     n_trees,
                     partial_sampling,
                     max_depth,
                     min_samples_leaf,
                     min_samples_split,
                     min_purity_increase;
                     rng = seed)

# And we generate predictions and measure their error
(ŷtrain,ŷval,ŷtest) = DecisionTree.apply_forest.([model],[xtrain,xval,xtest]);

# Let's benchmark the DecisionTrees.jl Random Forest training
@btime  DecisionTree.build_forest(ytrain, convert(Matrix,xtrain),
                     n_subfeatures,
                     n_trees,
                     partial_sampling,
                     max_depth,
                     min_samples_leaf,
                     min_samples_split,
                     min_purity_increase;
                     rng = seed);
# DecisionTrees.jl makes a good job in optimising the Random Forest algorithm, as it is over 3 times faster that BetaML.

(rmeTrain, rmeVal, rmeTest) = meanRelError.([ŷtrain,ŷval,ŷtest],[ytrain,yval,ytest],normRec=false)
#src 0.031235291992120055,0.17271800981298888,0.3141747075399198
# However the error on the test set remains relativly high. The very low error level on the training set is a sign that it overspecialised on the training set, and we should have better ran a dedicated hyper-parameter tuning function for the model.

@test rmeTest <= 0.36 #src

# Finally we plot the DecisionTree.jl predictions alongside the observed value:
ŷtrainfull = vcat(ŷtrain,fill(missing,nval+ntest))
ŷvalfull   = vcat(fill(missing,ntrain), ŷval, fill(missing,ntest))
ŷtestfull  = vcat(fill(missing,ntrain+nval), ŷtest)
plot(data[:,:dteday],[data[:,:cnt] ŷtrainfull ŷvalfull ŷtestfull], label=["obs" "train" "val" "test"], legend=:topleft, ylabel="daily rides", title="Daily bike sharing demand observed/estimated across the\n whole 2-years period (DT.jl RF)")

# Again, focusing on the testing data:
stc  = 620
endc = size(x,1)
plot(data[stc:endc,:dteday],[data[stc:endc,:cnt] ŷvalfull[stc:endc] ŷtestfull[stc:endc]], label=["obs" "val" "test"], legend=:bottomleft, ylabel="Daily rides", title="Focus on the testing period (DT.jl RF)")

# ### Conclusions of Decision Trees / Random Forests methods
# The error obtained employing DecisionTree.jl is significantly larger than those obtained using a BetaML random forest model, altought to be fair with DecisionTrees.jl we didn't tuned its hyper-parameters. Also, the DecisionTree.jl random forest model is much faster.
# This is partially due by the fact that, internally, DecisionTree.jl models optimise the algorithm by sorting the observations. BetaML trees/forests don't employ this optimisation and hence they can work with true categorical data for which ordering is not defined. An other explanation of this difference in speed is that BetaML Random Forest models accept `missing` values within the feature matrix.
# To sum up, BetaML random forests are ideal algorithms when we want to obtain good predictions in the most simpler way, even without tuning the hyper-parameters, and without spending time in cleaning ("munging") the feature matrix, as they accept almost "any kind" of data as it is.

# ## Neural Networks

# BetaML provides only _deep forward neural networks_, artificial neural network units where the individual "nodes" are arranged in _layers_, from the _input layer_, where each unit holds the input coordinate, through various _hidden layer_ transformations, until the actual _output_ of the model:

# ![Neural Networks](imgs/nn_scheme.png)

# In this layerwise computation, each unit in a particular layer takes input from _all_ the preceding layer units and it has its own parameters that are adjusted to perform the overall computation. The _training_ of the network consists in retrieving the coefficients that minimise a _loss_ function between the output of the model and the known data.
# In particular, a _deep_ (feedforward) neural network refers to a neural network that contains not only the input and output layers, but also hidden layers in between.

# Neural networks accept only numerical inputs. We hence need to convert all categorical data in numerical units. A common approach is to use the so-called "one-hot-encoding" where the catagorical values are converted into indicator variables (0/1), one for each possible value. This can be done in BetaML using the [`oneHotEncoder`](@ref) function:
seasonDummies  = convert(Array{Float64,2},oneHotEncoder(data[:,:season]))
weatherDummies = convert(Array{Float64,2},oneHotEncoder(data[:,:weathersit]))
wdayDummies    = convert(Array{Float64,2},oneHotEncoder(data[:,:weekday] .+ 1 ))

## We compose the feature matrix with the new dimensions obtained from the oneHotEncoder functions
x = hcat(Matrix{Float64}(data[:,[:instant,:yr,:mnth,:holiday,:workingday,:temp,:atemp,:hum,:windspeed]]),
         seasonDummies,
         weatherDummies,
         wdayDummies)
y = data[:,16];

# As usual, we split the data in training, validation and testing sets
((xtrain,xval,xtest),(ytrain,yval,ytest)) = partition([x,y],[0.75,0.125,1-0.75-0.125],shuffle=false)
(ntrain, nval, ntest) = size.([ytrain,yval,ytest],1)

# An other common operation with neural networks is to scale the feature vectors (X) and the labels (Y). The BetaML [`scale`](@ref) function, by default, scales the data such that each dimension has mean 0 and variance 1.

# Note that we can provide the function with different scale factors or specify the columns that shoudn't be scaled (e.g. those resulting from the one-hot encoding). Finally we can reverse the scaling (this is useful to retrieve the unscaled features from a model trained with scaled ones).

colsNotToScale = [2;4;5;10:23]
xScaleFactors   = getScaleFactors(xtrain,skip=colsNotToScale)
yScaleFactors   = ([0],[0.001]) # getScaleFactors(ytrain) # This just divide by 1000. Using full scaling of Y we may get negative demand.
xtrainScaled    = scale(xtrain,xScaleFactors)
xvalScaled      = scale(xval,xScaleFactors)
xtestScaled     = scale(xtest,xScaleFactors)
ytrainScaled    = scale(ytrain,yScaleFactors)
yvalScaled      = scale(yval,yScaleFactors)
ytestScaled     = scale(ytest,yScaleFactors)
D               = size(xtrain,2)

#-

# As we did above for decision trees and random forests, we select the best hyper-parameters by using the validation set.
# We consider here two hyper-parameters, the first one (`hiddenLayerSizeRange`) concerns the structure of the neural network, and in particular the size (in nodes) of the hidden layer, the second one (`epochRange`) concerns the training itself, and in particular the number of so-called "epochs" (number of iterations trough the whole dataset) to train the model.

# Again, we use here repetitions for simplicity, but a [cross-validation approach](@ref crossValidation) would have bee nmore appropriate:

function tuneHyperParameters(xtrain,ytrain,xval,yval;epochRange=50:50,hiddenLayerSizeRange=12:12,repetitions=5,rng=Random.GLOBAL_RNG)
    ## We start with an infinititly high error
    bestRme         = +Inf
    bestEpoch       = 0
    bestSize        = 0
    compLock        = ReentrantLock()

    ## Generate one random number generator per thread
    masterSeed = rand(rng,100:9999999999999) ## Some RNG have problems with very small seed. Also, the master seed has to be computed _before_ generateParallelRngs
    rngs       = generateParallelRngs(rng,Threads.nthreads())

    ## We loop over all possible hyperparameter combinations...
    parLengths = (length(epochRange),length(hiddenLayerSizeRange))
    Threads.@threads for ij in CartesianIndices(parLengths)
       (epoch,hiddenLayerSize)   = (epochRange[Tuple(ij)[1]], hiddenLayerSizeRange[Tuple(ij)[2]])
       tsrng = rngs[Threads.threadid()]
       joinedIndx = LinearIndices(parLengths)[ij]
       ## And here we make the seeding depending on the i of the loop, not the thread: hence we get the same results indipendently of the number of threads
       Random.seed!(tsrng,masterSeed+joinedIndx*10)
       totAttemptError = 0.0
       println("Testing epochs $epoch, layer size $hiddenLayerSize ...")
       ## We run several repetitions with the same hyperparameter combination to account for stochasticity...
       for r in 1:repetitions
           l1   = DenseLayer(D,hiddenLayerSize,f=relu,rng=tsrng) # Activation function is ReLU
           l2   = DenseLayer(hiddenLayerSize,hiddenLayerSize,f=identity,rng=tsrng)
           l3   = DenseLayer(hiddenLayerSize,1,f=relu,rng=tsrng)
           mynn = buildNetwork([l1,l2,l3],squaredCost,name="Bike sharing regression model") # Build the NN and use the squared cost (aka MSE) as error function
           ## Training it (default to ADAM)
           res  = train!(mynn,xtrain,ytrain,epochs=epoch,batchSize=8,optAlg=ADAM(),verbosity=NONE, rng=tsrng) # Use optAlg=SGD() to use Stochastic Gradient Descent
           ŷval = predict(mynn,xval)
           rmeVal  = meanRelError(ŷval,yval,normRec=false)
           totAttemptError += rmeVal
       end
       avgRme = totAttemptError / repetitions
       begin
           lock(compLock) ## This step can't be run in parallel...
           try
               ## Select this specific combination of hyperparameters if the error is the lowest
               if avgRme < bestRme
                 bestRme    = avgRme
                 bestEpoch  = epoch
                 bestSize   = hiddenLayerSize
               end
           finally
               unlock(compLock)
           end
       end
    end
    return (bestRme=bestRme,bestEpoch=bestEpoch,bestSize=bestSize)
end

#-

## Note: the following code block may take a bit to run...
epochsToTest     = [100,400]
hiddenLayerSizes = [5,15,30]
(bestRme,bestEpoch,bestSize) = tuneHyperParameters(xtrainScaled,ytrainScaled,xvalScaled,yvalScaled;epochRange=epochsToTest,hiddenLayerSizeRange=hiddenLayerSizes,repetitions=3,rng=copy(FIXEDRNG))
println("NN: $bestRme $bestEpoch $bestSize") #src
#src NN: 0.12631777965957725 100 5

# We can now build our feed-forward neaural network. We create three layers, the first layers will always have a input size equal to the dimensions of our data (the number of columns), and the output layer, for a simple regression where the predictions are scalars, it will always be one. The middle layer has the size we obtained from `tuneHyperParameters`.

# There are already several kind of layers available (and you can build your own kind by defining a new `struct` and implementing a few functions. See the [`Nn`](@ref nn_module) module documentation for details). Here we use only _dense_ layers, those found in typycal feed-fordward neural networks.

# For each layer, on top of its size (in "neurons") we can specify an _activation function_. Here we use the [`relu`](@ref) for the terminal layer (this will guarantee that our predictions are always positive) and `identity` for the hidden layer. Again, consult the `Nn` module documentation for other activation layers already defined, or use any function of your choice.

# Initial weight parameters can also be specified if needed. By default [`DenseLayer`](@ref) use the so-called _Xavier initialisation_.
l1   = DenseLayer(D,bestSize,f=relu,rng=copy(FIXEDRNG)) # Activation function is ReLU
l2   = DenseLayer(bestSize,bestSize,f=identity,rng=copy(FIXEDRNG))
l3   = DenseLayer(bestSize,1,f=relu,rng=copy(FIXEDRNG))

# Finally we "chain" the layers together and we assign a final loss function (agian, you can provide your own loss function, if those available in BetaML don't suit your needs) in order to create the "neural network" [`NN`](@ref) object with the [`buildNetwork`][@ref] function:

mynn = buildNetwork([l1,l2,l3],squaredCost,name="Bike sharing regression model") ## Build the NN and use the squared cost (aka MSE) as error function

# The above neural network will use automatic differentiation (using the [Zygote](https://github.com/FluxML/Zygote.jl) package) to compute the gradient used in the loss minimisation during the training step.
# It is also possible, for the layers that support it, to use manual differentiation. The network below is exactly equivalent to the one above, except it avoids automatic differentiation as we tell BetaML the functions to use as derivatives of the activation functions and of the overall network loss function:

mynnManual = buildNetwork([
        DenseLayer(D,bestSize,f=relu,df=drelu,rng=copy(FIXEDRNG)),
        DenseLayer(bestSize,bestSize,f=identity,df=didentity,rng=copy(FIXEDRNG)),
        DenseLayer(bestSize,1,f=relu,df=drelu,rng=copy(FIXEDRNG))
    ], squaredCost, name="Bike sharing regression model", dcf=dSquaredCost)

# We can now train the neural network with the best hyper-parameters using the function [`train!`](@ref). Note the esclamation point. By convention in julia, functions that end with an exclamation mark modify some of their inputs, normally the first one.

# Here `train!` has a question mark as indeed it modifies the neural network object, by updating its weights. If you want to keep a copy of the network before training (for example, if you want to train it in different independent ways), make a `deepcopy` of the `NN` object or first save its parameters with [`getParams`](@ref) and then re-apply the saved parameters with [`setParams!`](@ref).

# Several optimisation algorithms are available, and each accepts different parameters, like the _learning rate_ for the Stochastic Gradient Descent algorithm ([`SGD`](@ref], used by default) or the exponential decay rates for the  moments estimates for the [`ADAM`](@ref ) algorithm (that we use here, with the default parameters).
println("Final training of $bestEpoch epochs, with layer size $bestSize ...")
res  = train!(mynn,xtrainScaled,ytrainScaled,epochs=bestEpoch,batchSize=8,optAlg=ADAM(),rng=copy(FIXEDRNG)) ## Use optAlg=SGD() to use Stochastic Gradient Descent

# Let's benchmark the BetaML neural network training
@btime train!(mynnManual,xtrainScaled,ytrainScaled,epochs=bestEpoch,batchSize=8,optAlg=ADAM(),rng=copy(FIXEDRNG), verbosity=NONE);
# As we can see the model training is one order of magnitude slower than random forests, altought the memory requirement is approximatly the same.

#-

# To obtain the neural network predictions we apply the function `predict` to the feature matrix X for which we want to generate previsions, and then, in order to obtain the _unscaled_ estimates we use the [`scale`](@ref) function applied to the scaled values with the original scaling factors and the parameter `rev` set to `true`.
# Note the usage of the _pipe_ operator to avoid ugly `function1(function2(function3(...)))` nested calls:

ŷtrain = @pipe predict(mynn,xtrainScaled) |> scale(_, yScaleFactors,rev=true);
ŷval   = @pipe predict(mynn,xvalScaled)   |> scale(_, yScaleFactors,rev=true);
ŷtest  = @pipe predict(mynn,xtestScaled)  |> scale(_ ,yScaleFactors,rev=true);

#-
(mreTrain, mreVal, mreTest) = meanRelError.([ŷtrain,ŷval,ŷtest],[ytrain,yval,ytest],normRec=false)
#src 0.08839103036501561, 0.1482673840724729, 0.17613463255050987

# The error is much lower. Let's plot our predictions:
@test mreTest < 0.18 #src

# Again, we can start by plotting the estimated vs the observed value:
scatter(ytrain,ŷtrain,xlabel="daily rides",ylabel="est. daily rides",label=nothing,title="Est vs. obs in training period (NN)")
#-
scatter(yval,ŷval,xlabel="daily rides",ylabel="est. daily rides",label=nothing,title="Est vs. obs in validation period (NN)")
#-
scatter(ytest,ŷtest,xlabel="daily rides",ylabel="est. daily rides",label=nothing,title="Est vs. obs in testing period (NN)")
#-

# We now plot across the time dimension, first plotting the whole period (2 years):
ŷtrainfull = vcat(ŷtrain,fill(missing,nval+ntest))
ŷvalfull   = vcat(fill(missing,ntrain), ŷval, fill(missing,ntest))
ŷtestfull  = vcat(fill(missing,ntrain+nval), ŷtest)
plot(data[:,:dteday],[data[:,:cnt] ŷtrainfull ŷvalfull ŷtestfull], label=["obs" "train" "val" "test"], legend=:topleft, ylabel="daily rides", title="Daily bike sharing demand observed/estimated across the\n whole 2-years period  (NN)")

# ...and then focusing on the testing data
stc  = 620
endc = size(x,1)
plot(data[stc:endc,:dteday],[data[stc:endc,:cnt] ŷvalfull[stc:endc] ŷtestfull[stc:endc]], label=["obs" "val" "test"], legend=:bottomleft, ylabel="Daily rides", title="Focus on the testing period (NN)")

# ### Comparison with Flux

# We now apply the same Neural Network model using the [Flux](https://fluxml.ai/) framework, a dedicated neural network library, reusing the optimal parameters that we did learn in `tuneHyperParameters`

# We fix the default random number generator so that the Flux example gives a reproducible output
Random.seed!(123)

# We define the Flux neural network model and load it with data...
l1         = Flux.Dense(D,bestSize,Flux.relu)
l2         = Flux.Dense(bestSize,bestSize,identity)
l3         = Flux.Dense(bestSize,1,Flux.relu)
Flux_nn    = Flux.Chain(l1,l2,l3)
loss(x, y) = Flux.mse(Flux_nn(x), y)
ps         = Flux.params(Flux_nn)
nndata     = Flux.Data.DataLoader((xtrainScaled', ytrainScaled'), batchsize=8,shuffle=true)

Flux_nn2   = deepcopy(Flux_nn)      ## A copy for the time benchmarking
ps2        = Flux.params(Flux_nn2)  ## A copy for the time benchmarking

# We do the training of the Flux model...
Flux.@epochs bestEpoch Flux.train!(loss, ps, nndata, Flux.ADAM(0.001, (0.9, 0.8)));

# ..and we benchmark it..
@btime begin for i in 1:bestEpoch Flux.train!(loss, ps2, nndata, Flux.ADAM(0.001, (0.9, 0.8))) end end
#src # Quite surprisling, Flux training seems a bit slow. The actual results seems to depend from the actual hardware and by default Flux seems not to use multi-threading. While I suspect Flux scales better with larger networks and/or data, for these small examples on my laptop it is still a bit slower than BetaML even on a single thread.
# On this small example the speed of Flux is on the same order than BetaML (the actual difference seems to depend on the specific RNG seed and hardware), however I suspect that Flux scales much better with larger networks and/or data.

# We obtain the estimates...
ŷtrainf = @pipe Flux_nn(xtrainScaled')' |> scale(_,yScaleFactors,rev=true);
ŷvalf   = @pipe Flux_nn(xvalScaled')'   |> scale(_,yScaleFactors,rev=true);
ŷtestf  = @pipe Flux_nn(xtestScaled')'  |> scale(_,yScaleFactors,rev=true);

# ..and we compute the mean relative errors..
(mreTrain, mreVal, mreTest) = meanRelError.([ŷtrainf,ŷvalf,ŷtestf],[ytrain,yval,ytest],normRec=false)
#src 0.09808947560569008, 0.13646196170811473,0.16181699665523128
# .. finding an error not significantly different than the one obtained from BetaML.Nn.

#-
@test mreTest < 0.18 #src

# Plots:
scatter(ytrain,ŷtrainf,xlabel="daily rides",ylabel="est. daily rides",label=nothing,title="Est vs. obs in training period (Flux.NN)")
#-
scatter(yval,ŷvalf,xlabel="daily rides",ylabel="est. daily rides",label=nothing,title="Est vs. obs in validation period (Flux.NN)")
#-
scatter(ytest,ŷtestf,xlabel="daily rides",ylabel="est. daily rides",label=nothing,title="Est vs. obs in testing period (Flux.NN)")
#-
ŷtrainfullf = vcat(ŷtrainf,fill(missing,nval+ntest))
ŷvalfullf   = vcat(fill(missing,ntrain), ŷvalf, fill(missing,ntest))
ŷtestfullf  = vcat(fill(missing,ntrain+nval), ŷtestf)
plot(data[:,:dteday],[data[:,:cnt] ŷtrainfullf ŷvalfullf ŷtestfullf], label=["obs" "train" "val" "test"], legend=:topleft, ylabel="daily rides", title="Daily bike sharing demand observed/estimated across the\n whole 2-years period (Flux.NN)")
#-
stc = 620
endc = size(x,1)
plot(data[stc:endc,:dteday],[data[stc:endc,:cnt] ŷvalfullf[stc:endc] ŷtestfullf[stc:endc]], label=["obs" "val" "test"], legend=:bottomleft, ylabel="Daily rides", title="Focus on the testing period (Flux.NN)")

# ### Conclusions of Neural Network models

# If we strive for the most accurate predictions, deep neural networks are usually the best choice. However they are computationally expensive, so with limited resourses we may get better results by fine tuning and running many repetitions of "simpler" decision trees or even random forest models than a large naural network with insufficient hyper-parameter tuning.
# Also, we shoudl consider that decision trees/random forests are much simpler to work with.

# That said, specialised neural network libraries, like Flux, allow to use GPU and specialised hardware letting neural networks to scale with very large datasets.

# Still, for small and medium datasets, BetaML provides simpler yet customisable solutions that are accurate and fast.

# ## Summary

# This is the summary of the results we had trying to predict the daily bike sharing demand, given weather and calendar information of the day

# | Model                | Train rme     | Test rme|  Training time (ms)* | Training mem (MB)*  |
# |:-------------------- |:-------------:| --------:| ------------------- | ------------------- |
# | DT                   | 0.1266        | 0.2223   | 26.5                | 58                  |
# | RF                   | 0.0651        | 0.2223   | 362                 | 971                 |
# | RF (DecisionTree.jl) | 0.0312        | 0.3142   | 36                  | 11                  |
# | NN                   | 0.0884        | 0.1761   | 1768                | 758                 |
# | NN (Flux.jl)         | 0.0981        | 0.1618   | 1708                | 282                 |

# * on a Intel Core i5-8350U laptop

# Neural networks can be more precise than random forests models, but are more computationally expensive (and tricky to set up). When we compare BetaML with the algorithm-specific leading packages, we found similar results in terms of accuracy, but often the leading packages are better optimised and run more efficiently (but sometimes at the cost of being less versatile).

docs/src/tutorials/Regression - bike sharing/betaml_tutorial_regression_sharingBikes.jl

lessonsMaterial/00_INTRO/notebooks/Regression - bike sharing/betaml_tutorial_regression_sharingBikes.jl


# # Honeycomb Lattice in k-space

using LinearAlgebra
using Plots
using LatticeTools

simplifyname(s::AbstractString) = (
    s
    |> (x-> replace(x, r"(.+?)" => s"\1"))
    |> (x-> replace(x, r"(.+?)" => s"[\1]"))
)

function makewithin(extent::AbstractVector{<:Real})
    a, b, c, d = extent
    (x::Real, y::Real) -> ((a <= x <= b) && (c <= y <= d))
end
mkpath("example_little_symmetry_kspace_honeycomb")

# ## Honeycome Lattice

# ### Define UnitCell
latticevectors = [1 -0.5; 0 sqrt(3)*0.5];
unitcell = makeunitcell(latticevectors; SiteType=String);
addsite!(unitcell, "A", carte2fract(unitcell, [0.5, 0.5 / sqrt(3)]));
addsite!(unitcell, "B", carte2fract(unitcell, [0.5, -0.5 / sqrt(3)]));

# ### Plot Lattice
let bravais_lattice = [], a_sites = [], b_sites = []
    extent = [-2, 2, -2, 2]
    within = makewithin(extent)
    for i1 in -5:5, i2 in -5:5
        R = latticevectors * [i1, i2]
        within(R...) && push!(bravais_lattice, R)
        for (orb_name, orb_fc) in unitcell.sites
            orb_cc = fract2carte(unitcell, orb_fc)
            r = R + orb_cc
            within(r...) && push!(orb_name == "A" ? a_sites : b_sites, r)
        end
    end
    bravais_lattice = hcat(bravais_lattice...)
    a_sites = hcat(a_sites...)
    b_sites = hcat(b_sites...)

    img = plot(aspect_ratio=1)
    scatter!(bravais_lattice[1,:], bravais_lattice[2,:], color="black", markershape=:circle, markersize=3, label="Bravais")
    scatter!(a_sites[1,:], a_sites[2,:], color="red", markerstrokewidth=0, markersize=5, label="A")
    scatter!(b_sites[1,:], b_sites[2,:], color="blue", markerstrokewidth=0, markersize=5, label="B")
    xlims!(extent[1], extent[2])
    ylims!(extent[3], extent[4])
    savefig(img, "example_little_symmetry_kspace_honeycomb/realspace.svg")
end

# ![](example_little_symmetry_kspace_honeycomb/realspace.svg)

# ## √3 × √3 Supercell

# ### Make Superlattice
shape = [2 -1; 1 1]
lattice = makelattice(unitcell, shape)

# ### Plot Lattice
let bravais_lattice = [], a_sites = [], b_sites = []
    extent = [-4, 4, -4, 4]
    within = makewithin(extent)

    for i1 in -5:5, i2 in -5:5
        R = lattice.supercell.latticevectors * [i1, i2]
        within(R...) && push!(bravais_lattice, R)
        for (orb_name, orb_fc) in lattice.supercell.sites
            orb_cc = fract2carte(lattice.supercell, orb_fc)
            r = R + orb_cc
            within(r...) && push!(orb_name[1] == "A" ? a_sites : b_sites, r)
        end
    end
    bravais_lattice = hcat(bravais_lattice...)
    a_sites = hcat(a_sites...)
    b_sites = hcat(b_sites...)

    img = plot(aspect_ratio=1)
    scatter!(bravais_lattice[1,:], bravais_lattice[2,:], color="black", markershape=:circle, markersize=3, label="Bravais (supercell)")
    scatter!(a_sites[1,:], a_sites[2,:], color="red", markerstrokewidth=0, markersize=5, label="A")
    scatter!(b_sites[1,:], b_sites[2,:], color="blue", markerstrokewidth=0, markersize=5, label="B")
    xlims!(extent[1], extent[2])
    ylims!(extent[3], extent[4])
    savefig(img, "example_little_symmetry_kspace_honeycomb/realspace-root3xroot3.svg")
end

# ![](example_little_symmetry_kspace_honeycomb/realspace-root3xroot3.svg)

# ### Symmetries
tsym = FiniteTranslationSymmetry(lattice)
psym = project(PointSymmetryDatabase.get(25), [1 0 0; 0 1 0])
println("point symmetry (Hermann Mauguin): $(psym.hermann_mauguin)")
println("translation symmetry compatible with point symmetry? ", iscompatible(tsym, psym))

println("irreducible representations:")
for tsym_irrep_index in 1:num_irreps(tsym)
    kf = tsym.fractional_momenta[tsym_irrep_index]
    kc = lattice.unitcell.reducedreciprocallatticevectors * kf
    psym_little = little_symmetry(tsym, tsym_irrep_index, psym)
    print("- irrep index: $tsym_irrep_index\n")
    print("  fractional momentum: [")
    for (i, x) in enumerate(kf)
        i != 1 && print(", ")
        print(x)
    end
    print("]\n")
    print("  momentum: [")
    for (i, x) in enumerate(kc)
        i != 1 && print(", ")
        print(x)
    end
    print("]\n")
    print("  is point symmetry compatible with momentum: $(iscompatible(tsym, tsym_irrep_index, psym))\n")
    print("  little point symmetry: $(psym_little.hermann_mauguin)\n")
end

# ### Plot Momentum Space
let
    extent = [-1.1, 1.1, -1.4, 1.4]
    within = makewithin(extent)

    reciprocallatticepoints = []
    kpoints = []
    littlegroupnames = String[]
    generatornames = String[]
    kpointnames = String[]

    for tsym_irrep_index in 1:num_irreps(tsym)
        psym_little = little_symmetry(tsym, tsym_irrep_index, psym)
        kf = tsym.fractional_momenta[tsym_irrep_index]
        kc = lattice.unitcell.reducedreciprocallatticevectors * kf
        for i1 in -2:2, i2 in -2:2
            G = lattice.unitcell.reducedreciprocallatticevectors * [i1, i2]
            within(G...) && push!(reciprocallatticepoints, G)
            k = G + kc
            if within(k...)
                push!(kpoints, k)
                push!(littlegroupnames, psym_little.hermann_mauguin)
                push!(generatornames, join(simplifyname.(psym_little.element_names[psym_little.generators]), "\n"))
                push!(kpointnames, "$tsym_irrep_index")
            end
        end
    end
    reciprocallatticepoints = hcat(reciprocallatticepoints...)
    kpoints = hcat(kpoints...)

    img = plot(size=(400, 500), xlims=(extent[1]-0.1, extent[2]+0.1), ylims=(extent[3]-0.1, extent[4]+0.1), aspect_ratio=1)
    scatter!(reciprocallatticepoints[1,:], reciprocallatticepoints[2,:],
             markershape=:circle, markersize=6, markercolor=RGBA(1,1,1,0), markerstrokecolor=RGBA(1,0,0,1), label="reciprocal lattice")
    scatter!(kpoints[1,:], kpoints[2,:],
             markershape=:circle, markersize=4, markercolor=RGBA(0,0,1,0.5), markerstrokecolor=RGBA(0,0,1,0.5),
             series_annotations=[Plots.text("k[$x]", 8, :left, :bottom) for x in kpointnames], label="momentum points")
    scatter!(kpoints[1,:], kpoints[2,:],
             markershape=:circle, markersize=0, color=RGBA(0,0,0,0),
             series_annotations=[Plots.text("$x ", 8, :right, :top) for x in littlegroupnames], label="")
    scatter!(kpoints[1,:], kpoints[2,:],
             markershape=:circle, markersize=0, color=RGBA(0,0,0,0),
             series_annotations=[Plots.text("$x", 6, :left, :top) for x in generatornames], label="")
    savefig(img, "example_little_symmetry_kspace_honeycomb/momentumspace-root3xroot3.svg")
end

# ![](example_little_symmetry_kspace_honeycomb/momentumspace-root3xroot3.svg)

# ## 2√3 × 2√3 Supercell

# ### Make Superlattice
shape = [2 2; -2 4]
lattice = makelattice(unitcell, shape)

# ### Plot Lattice
let bravais_lattice = [], a_sites = [], b_sites = []
    extent = [-8, 8, -8, 8]
    within = makewithin(extent)

    for i1 in -5:5, i2 in -5:5
        R = lattice.supercell.latticevectors * [i1, i2]
        within(R...) && push!(bravais_lattice, R)
        for (orb_name, orb_fc) in lattice.supercell.sites
            orb_cc = fract2carte(lattice.supercell, orb_fc)
            r = R + orb_cc
            within(r...) && push!(orb_name[1] == "A" ? a_sites : b_sites, r)
        end
    end
    bravais_lattice = hcat(bravais_lattice...)
    a_sites = hcat(a_sites...)
    b_sites = hcat(b_sites...)

    img = plot(aspect_ratio=1)
    scatter!(bravais_lattice[1,:], bravais_lattice[2,:], color="black", markershape=:circle, markersize=5, label="Bravais (supercell)")
    scatter!(a_sites[1,:], a_sites[2,:], color="red", markerstrokewidth=0, markersize=3, label="A")
    scatter!(b_sites[1,:], b_sites[2,:], color="blue", markerstrokewidth=0, markersize=3, label="B")
    xlims!(extent[1], extent[2])
    ylims!(extent[3], extent[4])
    savefig(img, "example_little_symmetry_kspace_honeycomb/realspace-2root3x2root3.svg")
end

# ![](example_little_symmetry_kspace_honeycomb/realspace-2root3x2root3.svg)

# ### Symmetries
tsym = FiniteTranslationSymmetry(lattice)
psym = project(PointSymmetryDatabase.get(25), [1 0 0; 0 1 0])
println("point symmetry (Hermann Mauguin): $(psym.hermann_mauguin)")
println("translation symmetry compatible with point symmetry? ", iscompatible(tsym, psym))

println("irreducible representations:")
for tsym_irrep_index in 1:num_irreps(tsym)
    kf = tsym.fractional_momenta[tsym_irrep_index]
    kc = lattice.unitcell.reducedreciprocallatticevectors * kf
    psym_little = little_symmetry(tsym, tsym_irrep_index, psym)
    print("- irrep index: $tsym_irrep_index\n")
    print("  fractional momentum: [")
    for (i, x) in enumerate(kf)
        i != 1 && print(", ")
        print(x)
    end
    print("]\n")
    print("  momentum: [")
    for (i, x) in enumerate(kc)
        i != 1 && print(", ")
        print(x)
    end
    print("]\n")
    print("  is point symmetry compatible with momentum: $(iscompatible(tsym, tsym_irrep_index, psym))\n")
    print("  little point symmetry: $(psym_little.hermann_mauguin)\n")
end

# ### Plot Momentum Space
let
    extent = [-1.2, 1.2, -1.6, 1.6]
    extent = [-1.1, 1.1, -1.4, 1.4]
    within = makewithin(extent)

    reciprocallatticepoints = []
    kpoints = []
    littlegroupnames = String[]
    generatornames = String[]
    kpointnames = String[]

    for tsym_irrep_index in 1:num_irreps(tsym)
        psym_little = little_symmetry(tsym, tsym_irrep_index, psym)
        kf = tsym.fractional_momenta[tsym_irrep_index]
        kc = lattice.unitcell.reducedreciprocallatticevectors * kf
        for i1 in -2:2, i2 in -2:2
            G = lattice.unitcell.reducedreciprocallatticevectors * [i1, i2]
            within(G...) && push!(reciprocallatticepoints, G)
            k = G + kc
            if within(k...)
                push!(kpoints, k)
                push!(littlegroupnames, psym_little.hermann_mauguin)
                push!(generatornames, join(simplifyname.(psym_little.element_names[psym_little.generators]), "\n"))
                push!(kpointnames, "$tsym_irrep_index")
            end
        end
    end
    reciprocallatticepoints = hcat(reciprocallatticepoints...)
    kpoints = hcat(kpoints...)

    img = plot(size=(600, 750), xlims=(extent[1]-0.1, extent[2]+0.1), ylims=(extent[3]-0.1, extent[4]+0.1), aspect_ratio=1)
    scatter!(reciprocallatticepoints[1,:], reciprocallatticepoints[2,:],
             markershape=:circle, markersize=6, markercolor=RGBA(1,1,1,0), markerstrokecolor=RGBA(1,0,0,1), label="reciprocal lattice")
    scatter!(kpoints[1,:], kpoints[2,:],
             markershape=:circle, markersize=4, markercolor=RGBA(0,0,1,0.5), markerstrokecolor=RGBA(0,0,1,0.5),
             series_annotations=[Plots.text("k[$x]", 8, :left, :bottom) for x in kpointnames], label="momentum points")
    scatter!(kpoints[1,:], kpoints[2,:],
             markershape=:circle, markersize=0, color=RGBA(0,0,0,0),
             series_annotations=[Plots.text("$x ", 8, :right, :top) for x in littlegroupnames], label="")
    scatter!(kpoints[1,:], kpoints[2,:],
             markershape=:circle, markersize=0, color=RGBA(0,0,0,0),
             series_annotations=[Plots.text("$x", 6, :left, :top) for x in generatornames], label="")
    savefig(img, "example_little_symmetry_kspace_honeycomb/momentumspace-2root3x2root3.svg")
end

# ![](example_little_symmetry_kspace_honeycomb/momentumspace-2root3x2root3.svg)

# ## 6 × 6 Supercell

# ### Make Superlattice
shape = [6 0; 0 6]
lattice = makelattice(unitcell, shape)

# ### Plot Lattice
let bravais_lattice = [], a_sites = [], b_sites = []
    extent = [-8, 8, -8, 8]
    within = makewithin(extent)

    for i1 in -5:5, i2 in -5:5
        R = lattice.supercell.latticevectors * [i1, i2]
        within(R...) && push!(bravais_lattice, R)
        for (orb_name, orb_fc) in lattice.supercell.sites
            orb_cc = fract2carte(lattice.supercell, orb_fc)
            r = R + orb_cc
            within(r...) && push!(orb_name[1] == "A" ? a_sites : b_sites, r)
        end
    end
    bravais_lattice = hcat(bravais_lattice...)
    a_sites = hcat(a_sites...)
    b_sites = hcat(b_sites...)

    img = plot(aspect_ratio=1)
    scatter!(bravais_lattice[1,:], bravais_lattice[2,:], color="black", markershape=:circle, markersize=5, label="Bravais (supercell)")
    scatter!(a_sites[1,:], a_sites[2,:], color="red", markerstrokewidth=0, markersize=3, label="A")
    scatter!(b_sites[1,:], b_sites[2,:], color="blue", markerstrokewidth=0, markersize=3, label="B")
    xlims!(extent[1], extent[2])
    ylims!(extent[3], extent[4])
    savefig(img, "example_little_symmetry_kspace_honeycomb/realspace-6x6.svg")
end

# ![](example_little_symmetry_kspace_honeycomb/realspace-6x6.svg)

# ### Symmetries
tsym = FiniteTranslationSymmetry(lattice)
psym = project(PointSymmetryDatabase.get(25), [1 0 0; 0 1 0])
println("point symmetry (Hermann Mauguin): $(psym.hermann_mauguin)")
println("translation symmetry compatible with point symmetry? ", iscompatible(tsym, psym))

println("irreducible representations:")
for tsym_irrep_index in 1:num_irreps(tsym)
    kf = tsym.fractional_momenta[tsym_irrep_index]
    kc = lattice.unitcell.reducedreciprocallatticevectors * kf
    psym_little = little_symmetry(tsym, tsym_irrep_index, psym)
    print("- irrep index: $tsym_irrep_index\n")
    print("  fractional momentum: [")
    for (i, x) in enumerate(kf)
        i != 1 && print(", ")
        print(x)
    end
    print("]\n")
    print("  momentum: [")
    for (i, x) in enumerate(kc)
        i != 1 && print(", ")
        print(x)
    end
    print("]\n")
    print("  is point symmetry compatible with momentum: $(iscompatible(tsym, tsym_irrep_index, psym))\n")
    print("  little point symmetry: $(psym_little.hermann_mauguin)\n")
end

# ### Plot Momentum Space
let
    extent = [-1.2, 1.2, -1.6, 1.6]
    extent = [-1.1, 1.1, -1.4, 1.4]
    within = makewithin(extent)

    reciprocallatticepoints = []
    kpoints = []
    littlegroupnames = String[]
    generatornames = String[]
    kpointnames = String[]

    for tsym_irrep_index in 1:num_irreps(tsym)
        psym_little = little_symmetry(tsym, tsym_irrep_index, psym)
        kf = tsym.fractional_momenta[tsym_irrep_index]
        kc = lattice.unitcell.reducedreciprocallatticevectors * kf
        for i1 in -2:2, i2 in -2:2
            G = lattice.unitcell.reducedreciprocallatticevectors * [i1, i2]
            within(G...) && push!(reciprocallatticepoints, G)
            k = G + kc
            if within(k...)
                push!(kpoints, k)
                push!(littlegroupnames, psym_little.hermann_mauguin)
                push!(generatornames, join(simplifyname.(psym_little.element_names[psym_little.generators]), "\n"))
                push!(kpointnames, "$tsym_irrep_index")
            end
        end
    end
    reciprocallatticepoints = hcat(reciprocallatticepoints...)
    kpoints = hcat(kpoints...)

    img = plot(size=(800, 1000), xlims=(extent[1]-0.1, extent[2]+0.1), ylims=(extent[3]-0.1, extent[4]+0.1), aspect_ratio=1)
    scatter!(reciprocallatticepoints[1,:], reciprocallatticepoints[2,:],
             markershape=:circle, markersize=6, markercolor=RGBA(1,1,1,0), markerstrokecolor=RGBA(1,0,0,1), label="reciprocal lattice")
    scatter!(kpoints[1,:], kpoints[2,:],
             markershape=:circle, markersize=4, markercolor=RGBA(0,0,1,0.5), markerstrokecolor=RGBA(0,0,1,0.5),
             series_annotations=[Plots.text("k[$x]", 8, :left, :bottom) for x in kpointnames], label="momentum points")
    scatter!(kpoints[1,:], kpoints[2,:],
             markershape=:circle, markersize=0, color=RGBA(0,0,0,0),
             series_annotations=[Plots.text("$x ", 8, :right, :top) for x in littlegroupnames], label="")
    scatter!(kpoints[1,:], kpoints[2,:],
             markershape=:circle, markersize=0, color=RGBA(0,0,0,0),
             series_annotations=[Plots.text("$x", 6, :left, :top) for x in generatornames], label="")
    savefig(img, "example_little_symmetry_kspace_honeycomb/momentumspace-6x6.svg")
end

# ![](example_little_symmetry_kspace_honeycomb/momentumspace-6x6.svg)


@testset "tracked: accessors" begin

    # TODO: move to properties
    tape = Tape()
    r = rand()
    tr = tracked(tape, r)

    @test gettape(tr) == tape
    @test getid(tr) == -1
    @test getvalue(tr) == r
    setvalue!(tr, r + 1.0)
    @test getvalue(tr) == r + 1.0

    a = rand(3,4)
    ta = tracked(tape, a)

    @test gettape(ta) == tape
    @test getid(ta) == -1
    @test getvalue(ta) == a
    setvalue!(ta, a .+ 1)
    @test getvalue(ta) == a .+ 1

end

@testset "tracked: ops" begin

    tape = Tape()
    xv = rand(); yv = rand()
    x = tracked(tape, xv)
    y = tracked(tape, yv)
    z = 2x^2 + y
    @test getvalue(z) == 2xv^2 + yv
    @test length(tape) == 5

    tape = Tape()
    xv = rand(3,4); yv = rand(4,3)
    x = tracked(tape, xv); y = tracked(tape, yv)
    z = x * y .+ 1.0
    @test getvalue(z) == xv * yv .+ 1.0
    @test length(tape) == 3

    # tracked * untracked
    tape = Tape()
    x = tracked(tape, rand(3,3))
    y = rand(3,3)
    z = x * y
    @test length(tape) == 2
    @test z.val == x.val * y

end


using RobustLeastSquares
using BaseTestNext

Logging.configure(level=Logging.DEBUG) # Useful to check convergence

# Test solve()
@testset "RobustLeastSqures" begin

    @testset "Solvers" begin
        s1 = 20
        s2 = 10

        rng = srand(0)
        A = rand(rng,s1,s2)
        b = rand(rng,s1)
        weights = 1 + 0.2(rand(rng,s1)-0.5)

        sol_qr = RobustLeastSquares.solve(A,b,weights,:qr)
        sol_normal = RobustLeastSquares.solve(A,b,weights,:normal)
        sol_cg = RobustLeastSquares.solve(A,b,weights,:cg)

        @test norm(sol_qr - sol_normal) < 1e-10
        @test norm(sol_qr - sol_cg) < 1e-10
    end

    # Test the MEstimators
    @testset "MEstimators" begin
        width = 2.0
        r = [5.0]

        s1 = 20
        s2 = 10

        rng = srand(0)
        A = rand(rng,s1,s2)
        b = rand(rng,s1)
        weights = 1 + 0.2(rand(rng,s1)-0.5)

        @test RobustLeastSquares.estimator_rho(r,L2Estimator())        ≈ [12.5]
        @test RobustLeastSquares.estimator_psi(r,L2Estimator())        ≈ [5.0]
        @test RobustLeastSquares.estimator_weight(r,L2Estimator())     ≈ [1.0]
        @test RobustLeastSquares.estimator_sqrtweight(r,L2Estimator()) ≈ [1.0]

        @test RobustLeastSquares.estimator_rho(r,L1Estimator())        ≈ [5.0]
        @test RobustLeastSquares.estimator_psi(r,L1Estimator())        ≈ [1.0]
        @test RobustLeastSquares.estimator_weight(r,L1Estimator())     ≈ [0.20]
        @test RobustLeastSquares.estimator_sqrtweight(r,L1Estimator()) ≈ [0.4472135954999579]

        # problem here: width not implemented correctly for L1L2Estimator
        @test RobustLeastSquares.estimator_rho(r,L1L2Estimator(1.0))        ≈ [5.3484692283495345]
        @test RobustLeastSquares.estimator_psi(r,L1L2Estimator(1.0))        ≈ [1.3608276348795434]
        @test RobustLeastSquares.estimator_weight(r,L1L2Estimator(1.0))     ≈ [0.2721655269759087]
        @test RobustLeastSquares.estimator_sqrtweight(r,L1L2Estimator(1.0)) ≈ [0.5216948600244291]

        @test RobustLeastSquares.estimator_rho(r,HuberEstimator(width))        ≈ [8.0]
        @test RobustLeastSquares.estimator_psi(r,HuberEstimator(width))        ≈ [2.0]
        @test RobustLeastSquares.estimator_weight(r,HuberEstimator(width))     ≈ [0.4]
        @test RobustLeastSquares.estimator_sqrtweight(r,HuberEstimator(width)) ≈ [0.6324555320336759]

        @test RobustLeastSquares.estimator_rho(r,FairEstimator(width))        ≈ [4.988948126018528]
        @test RobustLeastSquares.estimator_psi(r,FairEstimator(width))        ≈ [1.4285714285714286]
        @test RobustLeastSquares.estimator_weight(r,FairEstimator(width))     ≈ [0.2857142857142857]
        @test RobustLeastSquares.estimator_sqrtweight(r,FairEstimator(width)) ≈ [0.5345224838248488]

        @test RobustLeastSquares.estimator_rho(r,CauchyEstimator(width))        ≈ [3.9620029377331667]
        @test RobustLeastSquares.estimator_psi(r,CauchyEstimator(width))        ≈ [0.6896551724137931]
        @test RobustLeastSquares.estimator_weight(r,CauchyEstimator(width))     ≈ [0.13793103448275862]
        @test RobustLeastSquares.estimator_sqrtweight(r,CauchyEstimator(width)) ≈ [0.3713906763541037]

        @test RobustLeastSquares.estimator_rho(r,GemanEstimator())        ≈ [0.4807692307692308]
        @test RobustLeastSquares.estimator_psi(r,GemanEstimator())        ≈ [0.0073964497041420114]
        @test RobustLeastSquares.estimator_weight(r,GemanEstimator())     ≈ [0.0014792899408284023]
        @test RobustLeastSquares.estimator_sqrtweight(r,GemanEstimator()) ≈ [0.038461538461538464]

        @test RobustLeastSquares.estimator_rho(r,WelschEstimator(width))        ≈ [1.9961390917275446]
        @test RobustLeastSquares.estimator_psi(r,WelschEstimator(width))        ≈ [0.009652270681138546]
        @test RobustLeastSquares.estimator_weight(r,WelschEstimator(width))     ≈ [0.0019304541362277093]
        @test RobustLeastSquares.estimator_sqrtweight(r,WelschEstimator(width)) ≈ [0.04393693362340742]

        @test RobustLeastSquares.estimator_rho(r,TukeyEstimator(width))        ≈ [0.6666666666666666]
        @test RobustLeastSquares.estimator_psi(r,TukeyEstimator(width))        ≈ [0.0]
        @test RobustLeastSquares.estimator_weight(r,TukeyEstimator(width))     ≈ [0.0]
        @test RobustLeastSquares.estimator_sqrtweight(r,TukeyEstimator(width)) ≈ [0.0]

        @test RobustLeastSquares.estimator_sqrtweight(b,MultiEstimator(CauchyEstimator(width)=>1:10))[1:10] ≈ RobustLeastSquares.estimator_sqrtweight(b[1:10],CauchyEstimator(width))
        @test RobustLeastSquares.estimator_sqrtweight(b,MultiEstimator(CauchyEstimator(width)=>1:10))[11:20] ≈ RobustLeastSquares.estimator_sqrtweight(b[11:20],L2Estimator(width))
    end

    @testset "Reweighted least squares" begin
        width = 2.0
        r = [5.0]

        s1 = 20
        s2 = 10

        rng = srand(0)
        A = rand(rng,s1,s2)
        b = rand(rng,s1)
        weights = 1 + 0.2(rand(rng,s1)-0.5)

        # Test the overall solver works...
        est = MultiEstimator(CauchyEstimator(width)=>1:10)
        sol2_qr, res, weights = reweighted_lsqr(A,b,est;method=:qr, n_iter=30, refit = false)
        sol2_normal, res, weights = reweighted_lsqr(A,b,est;method=:normal, n_iter=30, refit = false)
        sol2_cg, res, weights = reweighted_lsqr(A,b,est;method=:cg, n_iter=30, refit = false)

        @test norm(sol2_qr - sol2_normal) < 1e-7
        @test norm(sol2_qr - sol2_cg) < 1e-7

        sol3_qr, res, weights = reweighted_lsqr(A,b,est;method=:qr, n_iter=30, refit = true)
        sol3_normal, res, weights = reweighted_lsqr(A,b,est;method=:normal, n_iter=30, refit = true)
        sol3_cg, res, weights = reweighted_lsqr(A,b,est;method=:cg, n_iter=30, refit = true)

        @test norm(sol3_qr - sol3_normal) < 1e-7
        @test norm(sol3_qr - sol3_cg) < 1e-7
    end
end


@testset "∫ newcom" begin
    st = raw"""
        \newcommand{ \coma }[ 1]{hello #1}
        \newcommand{ \comb} [2 ]{\coma{#1}, goodbye #1, #2!}
        Then \comb{auth1}{auth2}.
        """
    @test isapproxstr(st |> conv,
            """

            Then hello auth1, goodbye auth1, auth2$(Markdown.htmlesc("!")).

""")
end

@testset "∫ math" begin
    st = raw"""
        \newcommand{\E}[1]{\mathbb E\left[#1\right]}
        \newcommand{\eqa}[1]{\begin{eqnarray}#1\end{eqnarray}}
        \newcommand{\R}{\mathbb R}
        Then something like
        \eqa{ \E{f(X)} \in \R &\text{if}& f:\R\maptso\R }
        """
    @test isapproxstr(st |> conv,
            """

            Then something like
            \\[\\begin{array}{c} \\mathbb E\\left[ f(X)\\right] \\in \\mathbb R &\\text{if}& f:\\mathbb R\\maptso\\mathbb R \\end{array}\\]

""")
end

@testset "∫ re-com" begin # see #36
    st = raw"""
        \newcommand{\com}[1]{⭒!#1⭒}
        \com{A}
        \newcommand{\com}[1]{◲!#1◲}
        \com{A}
        """
    steps = st |> explore_md_steps
    @test steps[:inter_md].inter_md == "\n ##FDINSERT## \n\n ##FDINSERT## \n"
    @test st |> conv == "⭒A⭒\n◲A◲"
end

@testset "∫ div" begin # see #74
    st = raw"""
        Etc and `~~~` but hey.
        @@dd but `x` and `y`? @@
        done
        """
    @test isapproxstr(st |> conv,
            """

            Etc and ~~~ but hey.

but x and y?

            done

""")
end

@testset "∫ code-l" begin
    st = raw"""
        Some code
        ```julia
        struct P
            x::Real
        end
        ```
        done
        """
    @test isapproxstr(st |> conv,
            """

            Some code

            struct P
                x::Real
            end

            done
""")
end

@testset "∫ math-br" begin # see #73
    st = raw"""
        \newcommand{\R}{\mathbb R}
        $$
            \min_{x\in \R^n} \quad f(x)+i_C(x).
        $$
        """
    @test isapproxstr(st |> conv,
            """\\[ \\min_{x\\in \\mathbb R^n} \\quad f(x)+i_C(x). \\]""")
end

@testset "∫ insert" begin # see also #65
    st = raw"""
        \newcommand{\com}[1]{⭒!#1⭒}
        abc\com{A}\com{B}def.
        """
    @test st |> conv == "
abc⭒A⭒⭒B⭒def.
\n"

    st = raw"""
        \newcommand{\com}[1]{⭒!#1⭒}
        abc

        \com{A}\com{B}

        def.
        """
    @test st |> conv == "
abc
\n⭒A⭒⭒B⭒\n
def.
\n"

    st = raw"""
        \newcommand{\com}[1]{⭒!#1⭒}
        abc

        \com{A}

        def.
        """
    @test st |> conv == "
abc
\n⭒A⭒\n
def.
\n"

    st = raw"""
        \newcommand{\com}[1]{†!#1†}
        blah \com{a}
        * \com{aaa} tt
        * ss \com{bbb}
        """
    @test isapproxstr(st |> conv,
            """
blah †a†

†aaa† tt

ss †bbb†

""")
end

@testset "∫ br-rge" begin # see #70
    st = raw"""
           \newcommand{\scal}[1]{\left\langle#1\right\rangle}
           \newcommand{\E}{\mathbb E}
           exhibit A
           $\scal{\mu, \nu} = \E[X]$
           exhibit B
           $\E[X] = \scal{\mu, \nu}$
           end."""
    @test isapproxstr(st |> conv,
            """
exhibit A
            \\(\\left\\langle \\mu, \\nu\\right\\rangle = \\mathbb E[X]\\)
            exhibit B
            \\(\\mathbb E[X] = \\left\\langle \\mu, \\nu\\right\\rangle\\)
            end.
""")
end

@testset "∫ cond" begin
    stv(var1, var2) = """
        @def var1 = $var1
        @def var2 = $var2
        start
        ~~~
        {{ if var1 }} targ1 {{ else if var2 }} targ2 {{ else }} targ3 {{ end }}
        ~~~
        done
        """
    @test stv(true, true)   |> conv == "
start \n targ1 \n done
\n"
    @test stv(false, true)  |> conv == "
start \n targ2 \n done
\n"
    @test stv(false, false) |> conv == "
start \n targ3 \n done
\n"
end

@testset "∫ recurs" begin # see #97
    st = raw"""
        | A | B |
        | :---: | :---: |
        | C | D |
        """

    if VERSION >= v"1.4.0-"
        @test isapproxstr(st |> conv,
            """

A	B

C	D

""")
    else
        @test isapproxstr(st |> conv,
            """

A	B

C	D

""")
    end

    @test F.convert_md(st, isrecursive=true) |> chtml == st |> conv

    st = raw"""
        @@emptydiv
        @@emptycore
        @@
        @@
        """
    st |> conv == "

\n
\n"
end

@testset "∫ label" begin
    st = raw"""
        Blah blah
        ## some title
        and then an anchor here \label{anchor} done.
        """
    F.def_GLOBAL_LXDEFS!()
    r = st |> conv
    @test occursin("here  done.", r)
end


# Simplicial Complex

"""Simplicial Complex Type

Simplex parameter type `P` must have implemented total order function `isless`.
"""
mutable struct SimplicialComplex{P} <: AbstractComplex
    cells::Dict{Int,Vector{Simplex{P}}}   # cells per dimension
    # order::Dict{Pair{Int,Int},Int}      # order of cells
end
Base.show(io::IO, cplx::SimplicialComplex) = print(io, "SimplicialComplex($(size(cplx)))")
Base.copy(cplx::SimplicialComplex) = SimplicialComplex(deepcopy(cplx.cells))

# ---------------
# Private Methods
# ---------------

"""Add a simplex to the complex (as is) and return a complex size, a dimension of simplex and its index in it

This function **doesn't** add missing faces of the simplex to the complex. Use `addsimplex!` function for instead.
"""
function addsimplex(cplx::SimplicialComplex{P}, splx::Simplex{P}) where {P}
    d = dim(splx)
    d < 0 && return (0,d,0)
    !haskey(cplx.cells, d) && setindex!(cplx.cells, Simplex{P}[], d)
    i = length(cplx.cells[d])+1
    splx[:index] = i
    push!(cplx.cells[d], splx)
    # cplx.order[d=>i] = length(cplx.order)+1
    return splx
end

"""Add a simplex to the complex and all of its faces recursivly and return a complex size, a dimension of simplex and its index in it"""
function addsimplex!(cplx::SimplicialComplex{P}, splx::Simplex{P}) where {P}
    # cache already
    added = Set{Simplex{P}}()

    toprocess = Simplex{P}[]
    ret = Simplex{P}[]
    addedidxs = NTuple{3,Int}[]

    # add simplex to complex
    s = addsimplex(cplx, splx)
    push!(addedidxs, (sum(size(cplx)), dim(s), s[:index]))
    push!(ret, s)
    for f in faces(splx) # add simplex faces for processing
        push!(toprocess, f)
    end

    # process simplex faces and add missing simplexes
    while(length(toprocess) > 0)
        tmp = pop!(toprocess) # processing simplex faces
        d = dim(tmp)

        cplx[tmp, d] <= size(cplx, d) && continue # skip if already in complex
        tmp in added && continue # skip if already processed

        s = addsimplex(cplx, tmp) # add simples to the complex
        push!(addedidxs, (sum(size(cplx)), dim(s), s[:index]))
        push!(added, tmp) # mark as processed
        push!(ret, s)

        dim(tmp) == 0 && continue # do not create faces for 0-simplexes
        for f in faces(tmp) # add simplex faces for processing
            push!(toprocess, f)
        end
    end

    return ret
end

#
# Public Interface
#

celltype(cplx::SimplicialComplex{P}) where {P} = Simplex{P}

function cells(cplx::SimplicialComplex)
    CCT = valtype(cplx.cells)
    length(cplx.cells) == 0 && return CCT[] # no cell in complex
    dims = maximum(keys(cplx.cells))
    CCT[haskey(cplx.cells, d) ? cplx.cells[d] : CCT() for d in 0:dims]
end

cells(cplx::SimplicialComplex{P}, d::Int) where {P} = get(cplx.cells, d,  nothing)

dim(cplx::SimplicialComplex) = length(size(cplx))-1

function boundary(cplx::SimplicialComplex, idx::Int, d::Int, ::Type{R}) where {R}
    ch = Chain(d-1, R)
    if d == 0
        setdim!(ch, -1)
        return ch
    end
    pos = one(R)
    neg = -one(R)

    splx = cplx[idx, d]
    splx === nothing && return ch

    sgn = true
    for face in faces(splx)
        fidx = cplx[face, d-1]
        push!(ch, (sgn ? pos : neg), fidx)
        sgn = !sgn
    end

    return ch
end

function coboundary(cplx::SimplicialComplex, idx::Int, d::Int, ::Type{R}) where {R}

    d == 0 && return Chain(d+1, R)

    cbd = Dict{Int,Chain{R}}()
    for i in 1:size(cplx,d-1)
        cbd[i] = Chain(d+1, R)
    end

    for i in 1:size(cplx,d)
        bd = boundary(cplx, i, d, R)
        for (c, el) in bd
            push!(cbd[el], c, i)
        end
    end

    return cbd[idx]
end

Base.push!(cplx::SimplicialComplex{P}, splx::Simplex{P}; recursive=false) where {P} =
    recursive ? addsimplex!(cplx, splx) : [addsimplex(cplx, splx)]

#
# Constructors
#

SimplicialComplex(::Type{P}) where P = SimplicialComplex(Dict{Int,Vector{Simplex{P}}}())
(::Type{SimplicialComplex{P}})() where P = SimplicialComplex(P)

function SimplicialComplex(splxs::Simplex{P}...) where {P}
    cplx = SimplicialComplex(P)
    added = Set{Simplex}()
    toprocess = Vector{Simplex}()
    for splx in splxs
        push!(toprocess, splx) # add simplex to process its faces
        while(length(toprocess) > 0)
            tmp = pop!(toprocess) # processing simplex faces
            tmp in added && continue # skip if already processed
            addsimplex(cplx, tmp) # add simples to complex
            push!(added, tmp) # mark as processed
            for f in faces(tmp) # add simplex faces for processing
                push!(toprocess, f)
            end
        end
    end

    return cplx
end

SimplicialComplex(splxs::Vector{Simplex}) = SimplicialComplex(splxs...)

#
# Miscellaneous
#

function Base.read(io::IO, ::Type{SimplicialComplex{P}}) where {P}
    splxs = Simplex[]
    while !eof(io)
        l = chomp(readline(io))
        si = map(e->parse(P, e), split(l, ' '))
        push!(splxs, Simplex(si...))
    end
    return SimplicialComplex(splxs)
end

function Base.write(io::IO, cplx::SimplicialComplex)
    zsplxs = Dict{Int,Int}()
    for s in cplx.cells[0]
        zsplxs[s.vs[]] = s.idx
        write(io, "$(s.idx)")
        write(io, 0x0A)
    end
    for d in 1:dim(cplx)
        for s in cplx.cells[d]
            for (j,i) in enumerate(s.vs)
                write(io, "$(zsplxs[i])")
                if j>d
                    write(io, 0x0A)
                else
                    write(io, ' ')
                end
            end
        end
    end
    return
end


module SquarenessBenchmark

import Squareness
using Squareness: issquare_with_isqrt
using BenchmarkTools

issquare_with_padic_newton(n::Signed) = n < 0 ? false : Squareness.issquare_with_padic_newton(n % Unsigned)
# issquare_with_padic_newton(n::Unsigned) = Squareness.issquare_with_padic_newton(n)
issquare_with_padic_newton(n::BigInt) = n < 0 ? false : Squareness.issquare_with_padic_newton(n)

function _issquare_with_simple_newton(x::I) where {I <: Integer}
    x == 0 && return true
    xtz = trailing_zeros(x)
    isodd(xtz) && return false
    x >>= unsigned(xtz)
    o_xrt::I = x
    xrt::I = (x + 0x1) >> 0x1
    xrt == 0 && return false  # overflow
    while xrt < o_xrt
        o_xrt = xrt
        xrt = (xrt + x ÷ xrt) >> 0x1
    end
    xrt * xrt == x
end
issquare_with_simple_newton(x::Signed) = x < 0 ? false : _issquare_with_simple_newton(x % Unsigned)
# issquare_with_simple_newton(x::Unsigned) = _issquare_with_simple_newton(x)
issquare_with_simple_newton(x::BigInt) = x < 0 ? false : _issquare_with_simple_newton(x)

println("## Int64 ≥ 0\n")

rand_int64() = rand(Int64) & typemax(Int64)

println("### P-adic Newton Method\n")
println("```")
display(@benchmark issquare_with_padic_newton(n) setup=(n=rand_int64()))
println("\n```\n\n")

println("### With `isqrt()`\n")
println("```")
display(@benchmark issquare_with_isqrt(n) setup=(n=rand_int64()))
println("\n```\n\n")

println("### Simple Newton Method\n")
println("```")
display(@benchmark issquare_with_simple_newton(n) setup=(n=rand_int64()))
println("\n```\n\n")

println("## Int64 ≥ 0 × 100\n")

rand_int64(N) = rand(Int64, N) .& typemax(Int64)

println("### P-adic Newton Method\n")
println("```")
display(@benchmark issquare_with_padic_newton.(n) setup=(n=rand_int64(100)))
println("\n```\n\n")

println("### With `isqrt()`\n")
println("```")
display(@benchmark issquare_with_isqrt.(n) setup=(n=rand_int64(100)))
println("\n```\n\n")

println("### Simple Newton Method\n")
println("```")
display(@benchmark issquare_with_simple_newton.(n) setup=(n=rand_int64(100)))
println("\n```\n\n")

println("## Int32 ≥ 0\n")

rand_int32() = rand(Int32) & typemax(Int32)

println("### P-adic Newton Method\n")
println("```")
display(@benchmark issquare_with_padic_newton(n) setup=(n=rand_int32()))
println("\n```\n\n")

println("### With `isqrt()`\n")
println("```")
display(@benchmark issquare_with_isqrt(n) setup=(n=rand_int32()))
println("\n```\n\n")

println("### Simple Newton Method\n")
println("```")
display(@benchmark issquare_with_simple_newton(n) setup=(n=rand_int32()))
println("\n```\n\n")

println("## Int32 ≥ 0 × 100\n")

rand_int32(N) = rand(Int32, N) .& typemax(Int32)

println("### P-adic Newton Method\n")
println("```")
display(@benchmark issquare_with_padic_newton.(n) setup=(n=rand_int32(100)))
println("\n```\n\n")

println("### With `isqrt()`\n")
println("```")
display(@benchmark issquare_with_isqrt.(n) setup=(n=rand_int32(100)))
println("\n```\n\n")

println("### Simple Newton Method\n")
println("```")
display(@benchmark issquare_with_simple_newton.(n) setup=(n=rand_int32(100)))
println("\n```\n\n")

println("## Int16 ≥ 0\n")

rand_int16() = rand(Int16) & typemax(Int16)

println("### P-adic Newton Method\n")
println("```")
display(@benchmark issquare_with_padic_newton(n) setup=(n=rand_int16()))
println("\n```\n\n")

println("### With `isqrt()`\n")
println("```")
display(@benchmark issquare_with_isqrt(n) setup=(n=rand_int16()))
println("\n```\n\n")

println("### Simple Newton Method\n")
println("```")
display(@benchmark issquare_with_simple_newton(n) setup=(n=rand_int16()))
println("\n```\n\n")

println("## Int16 ≥ 0 × 100\n")

rand_int16(N) = rand(Int16, N) .& typemax(Int16)

println("### P-adic Newton Method\n")
println("```")
display(@benchmark issquare_with_padic_newton.(n) setup=(n=rand_int16(100)))
println("\n```\n\n")

println("### With `isqrt()`\n")
println("```")
display(@benchmark issquare_with_isqrt.(n) setup=(n=rand_int16(100)))
println("\n```\n\n")

println("### Simple Newton Method\n")
println("```")
display(@benchmark issquare_with_simple_newton.(n) setup=(n=rand_int16(100)))
println("\n```\n\n")

println("## Int8 ≥ 0\n")

rand_int8() = rand(Int8) & typemax(Int8)

println("### P-adic Newton Method\n")
println("```")
display(@benchmark issquare_with_padic_newton(n) setup=(n=rand_int8()))
println("\n```\n\n")

println("### With `isqrt()`\n")
println("```")
display(@benchmark issquare_with_isqrt(n) setup=(n=rand_int8()))
println("\n```\n\n")

println("### Simple Newton Method\n")
println("```")
display(@benchmark issquare_with_simple_newton(n) setup=(n=rand_int8()))
println("\n```\n\n")

println("## Int8 ≥ 0 × 100\n")

rand_int8(N) = rand(Int8, N) .& typemax(Int8)

println("### P-adic Newton Method\n")
println("```")
display(@benchmark issquare_with_padic_newton.(n) setup=(n=rand_int8(100)))
println("\n```\n\n")

println("### With `isqrt()`\n")
println("```")
display(@benchmark issquare_with_isqrt.(n) setup=(n=rand_int8(100)))
println("\n```\n\n")

println("### Simple Newton Method\n")
println("```")
display(@benchmark issquare_with_simple_newton.(n) setup=(n=rand_int8(100)))
println("\n```\n\n")

println("## Int128 ≥ 0\n")

rand_int128() = rand(Int128) & typemax(Int128)

println("### P-adic Newton Method\n")
println("```")
display(@benchmark issquare_with_padic_newton(n) setup=(n=rand_int128()))
println("\n```\n\n")

println("### With `isqrt()`\n")
println("```")
display(@benchmark issquare_with_isqrt(n) setup=(n=rand_int128()))
println("\n```\n\n")

println("### Simple Newton Method\n")
println("```")
display(@benchmark issquare_with_simple_newton(n) setup=(n=rand_int128()))
println("\n```\n\n")

println("## Int128 ≥ 0 × 100\n")

rand_int128(N) = rand(Int128, N) .& typemax(Int128)

println("### P-adic Newton Method\n")
println("```")
display(@benchmark issquare_with_padic_newton.(n) setup=(n=rand_int128(100)))
println("\n```\n\n")

println("### With `isqrt()`\n")
println("```")
display(@benchmark issquare_with_isqrt.(n) setup=(n=rand_int128(100)))
println("\n```\n\n")

println("### Simple Newton Method\n")
println("```")
display(@benchmark issquare_with_simple_newton.(n) setup=(n=rand_int128(100)))
println("\n```\n\n")

println("## BigInt ≥ 0\n")

rand_mpz() = (big(rand(UInt128)) << rand(0:128)) + rand(UInt128)

println("### P-adic Newton Method\n")
println("```")
display(@benchmark issquare_with_padic_newton(n) setup=(n=rand_mpz()))
println("\n```\n\n")

println("### With `isqrt()`\n")
println("```")
display(@benchmark issquare_with_isqrt(n) setup=(n=rand_mpz()))
println("\n```\n\n")

println("### Simple Newton Method\n")
println("```")
display(@benchmark issquare_with_simple_newton(n) setup=(n=rand_mpz()))
println("\n```\n\n")

end


using Test, Random, FillArrays
import LuxurySparse: IMatrix, PermMatrix

Random.seed!(2)

p1 = IMatrix{4}()
sp = sprand(ComplexF64, 4,4, 0.5)
ds = rand(ComplexF64, 4,4)
pm = PermMatrix([2,3,4,1], randn(4))
v = [0.5, 0.3im, 0.2, 1.0]
dv = Diagonal(v)

@testset "basic" begin
    @test p1==copy(p1)
    @test eltype(p1) == Bool
    @test size(p1) == (4, 4)
    @test size(p1, 1) == size(p1, 2) == 4
    @test Matrix(p1) == [1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1]
end

@testset "conversion" begin
    for mat in [p1, pm, dv]
        @test mat == SparseMatrixCSC(mat)
        @test mat == Matrix(mat)
    end
    for mat in [p1, pm, dv]
        @test mat == PermMatrix(mat)
    end
    @test Diagonal(p1) == p1

    @test SparseMatrixCSC(Diagonal(Fill(2, 4))) ≈ Diagonal(Fill(2, 4))
end

@testset "sparse" begin
    @test nnz(p1) == 4
    @test nonzeros(p1) == ones(4)
end

@testset "linalg" begin
    for op in [conj, real, transpose, copy, inv]
        @test op(p1) == Matrix(I, 4, 4)
        @test typeof(op(p1)) == typeof(p1)
    end
    @test imag(p1) == zeros(4, 4)
    @test p1' == Matrix(I, 4, 4)

    # This will be lazy evaluated in 0.7+
    @static if VERSION < v"0.7-"
        @test typeof(p1') == typeof(p1)
    end

    @test ishermitian(p1)
end

@testset "elementary" begin
    @test all(isapprox.(conj(p1), conj(Matrix(p1))))
    @test all(isapprox.(real(p1), real(Matrix(p1))))
    @test all(isapprox.(imag(p1), imag(Matrix(p1))))
end

@testset "basicmath" begin
    @test p1*2im == Matrix(p1)*2im
    @test p1/2.0 == Matrix(p1)/2.0
end


@testset "$TEST $G" begin

g = CompleteGraph(5, G)
klist = cores(g)
@test klist == fill(4, 5)

g = PathGraph(3)
klist = cores(g)
@test klist == fill(1, 3)

g = CycleGraph(3)
klist = cores(g)
@test klist == fill(2, 3)

add_vertex!(g)
add_edge!(g, 3, 4)
klist = cores(g)
@test klist == [2, 2, 2, 1]

h, vmap = kcore(g, 2)
@test h == CycleGraph(3)
@test length(vmap) == 0
@test vmap == VertexMap(g, vertextype(g))
end # testset


################################################################################
####################### Approximation error plot ###############################
################################################################################
function approx_error2(DVEC::Array{Array{Float64,1},1},
                       T::Vector{String}, L::Vector{Tuple{Symbol,Symbol}},
                       frac::Float64 = 0.30)
    # This version plots the relative L2 errors against
    # the FRACTION of coefficients retained.
    plot(xaxis = "Fraction of Coefficients Retained",
         yaxis = "Relative Approximation Error")
    for i = 1:length(DVEC)
        dvec = DVEC[i]
        N = length(dvec)
        dvec_norm = norm(dvec,2)
        dvec_sort = sort(dvec.^2) # the smallest first
        er = sqrt.(reverse(cumsum(dvec_sort)))/dvec_norm
        er[er .== 0.0] .= minimum(er[er .!= 0.0]) # avoid blowup by taking log in the plot below
        # er is the relative L^2 error of the whole thing: length(er)=N.
        p = Int64(floor(frac*N)) + 1 # upper limit
        plot!(frac*(0:(p-1))/(p-1), er[1:p], yaxis=:log, xlims = (0.,frac),
              label = T[i], line = L[i], linewidth = 2, grid = false)
    end
    display(current())
end

function approx_error3(DVEC::Array{Array{Float64,1},1},
                       T::Vector{String}, L::Vector{Tuple{Symbol,Symbol}},
                       N::Int64)
    # This version plots the relative L2 errors against
    # the NUMBER of coefficients retained.
    plot(xaxis = "Number of Coefficients Retained",
         yaxis = "Relative Approximation Error")
    for i = 1:length(DVEC)
        dvec = DVEC[i]
        dvec_norm = norm(dvec,2)
        dvec_sort = sort(dvec.^2) # the smallest first
        er = sqrt.(reverse(cumsum(dvec_sort)))/dvec_norm
        er[er .== 0.0] .= minimum(er[er .!= 0.0]) # avoid blowup by taking log in the plot below
        # er is the relative L^2 error of the whole thing: length(er)=N.
        plot!(0:N-1, er[1:N], yaxis=:log, label = T[i], line = L[i],
              linewidth = 2, grid = false)
    end
    display(current())
end

################################################################################
############ Computing a weight matrix using the Gaussian affinity #############
################################################################################
function image_Gaussian_affinity(img::Matrix{Float64}, r::Int64, σ::Float64)
# Get the neighbors for affinity matrix computation
    l = 2*r + 1
    temp_x, temp_y = fill(1,l^2), fill(1,l^2)
    temp_xy = CartesianIndices((1:l,1:l))
    for i = 1:l^2
        temp_x[i], temp_y[i] = temp_xy[i][1], temp_xy[i][2]
    end
    # (r+1, r+1) is the index of the center location.
    temp_ind = ((temp_x .- (r + 1)).^2 + (temp_y .- (r + 1)).^2) .<= r^2
    # Now, temp_ind indicates those points withinin the circle of radius r
    # from the center while neighbor_x, neighbor_y represent relative positions
    # of those points from the center.
    neighbor_x = temp_x[temp_ind] .- (r + 1)
    neighbor_y = temp_y[temp_ind] .- (r + 1)
   # So, for any index (x, y), points within (x ± neighbor_x, y ± neighbor_y) are
   # its neighbors for the purpose of calculating affinity

# Create affinity matrix
    m, n = size(img)
    sig = img[:]
    W = fill(0., (m*n,m*n))
    for i = 1:m*n
        cur = CartesianIndices((m,n))[i]
        for j = 1:length(neighbor_x) # assuming dim(neighbor_x) == dim(neighbor_y)
            if 1 <= cur[1] + neighbor_x[j] <= m && 1 <= cur[2] + neighbor_y[j] <= n
                tempd = LinearIndices((m,n))[cur[1] + neighbor_x[j], cur[2] + neighbor_y[j]]
                W[i,tempd] = exp(-(sig[i] - sig[tempd])^2/σ)
            end
        end
    end
    return sparse(W)
end # end of the image_Gaussian_affinity function

################################################################################
######## Display top 9 basis vectors of various bases for an image data ########
################################################################################
function top_vectors_plot2(dvec::Array{Float64, 2}, m::Int64, n::Int64, BS::BasisSpec, GP::GraphPart;
                           clims::Tuple{Float64,Float64} = (-0.01, 0.01))

    # Get the indices from max to min in terms of absolute value of the coefs
    # Note that m*n == length(dvec); m and n are the image size.
    sorted_ind = sortperm(abs.(dvec[:]), rev = true);

    # Set up the layout as 3 x 3 subplots
    plot(9, layout = (3,3), framestyle = :none, legend = false)

    # Do the display!
    for i=1:9
        dvecT = fill(0., size(dvec))
        dvecT[sorted_ind[i]] = 1
        f = ghwt_synthesis(dvecT, GP, BS)
        heatmap!(reshape(f, (m,n)), subplot=i, ratio=1, yaxis=:flip,
                 showaxis=false, ticks = false, color = :grays, clims = clims)
    end
    display(current())
end
################################################################################

function top_vectors_plot3(dvec::Array{Float64, 2}, m::Int64, n::Int64, BS::BasisSpec,
    GP::GraphPart; clims::Tuple{Float64,Float64} = (-0.01, 0.01), K::Int64 = 9)

    # Get the indices from max to min in terms of absolute value of the coefs
    # Note that m*n == length(dvec); m and n are the image size.
    sorted_ind = sortperm(abs.(dvec[:]), rev = true);

    # Set up the layout as 3 x 3 subplots
    plot(K, layout = (Int(sqrt(K)), Int(sqrt(K))), framestyle = :none, legend = false)

    # Do the display!
    for i=1:K
        dvecT = fill(0., size(dvec))
        dvecT[sorted_ind[i]] = 1
        f = ghwt_synthesis(dvecT, GP, BS)
        heatmap!(reshape(f, (m,n)), subplot=i, ratio=1, yaxis=:flip,
                 showaxis=false, ticks = false, color = :grays, clims = clims)
    end
    display(current())
end
################################################################################


struct Solution
    ret_val::Int
    primal::Vector{Float64}
    dual_eq::Vector{Float64}
    dual_ineq::Vector{Float64}
    slack::Vector{Float64}
    objval::Float64
end
Solution() = Solution(0,
                      Float64[], Float64[], Float64[], Float64[], NaN)

mutable struct Data
    m::Int
    n::Int
    IA::Vector{Int}
    JA::Vector{Int}
    VA::Vector{Float64}
    b::Vector{Float64}
    IG::Vector{Int}
    JG::Vector{Int}
    VG::Vector{Float64}
    h::Vector{Float64}
    objconstant::Float64
    c::Vector{Float64}
end

mutable struct Cone
    f::Int # number of linear equality constraints
    l::Int # length of LP cone
    q::Int # length of SOC cone
    qa::Vector{Int} # array of second-order cone constraints
    ep::Int # number of primal exponential cone triples
    eqsetconstant::Dict{Int, Float64} # For the constant of EqualTo
    eqnrows::Dict{Int, Int} # The number of rows of Zeros
    ineqsetconstant::Dict{Int, Float64} # For the constant of LessThan and GreaterThan
    ineqnrows::Dict{Int, Int} # The number of rows of each vector sets except Zeros
    function Cone()
        new(0, 0, 0, Int[], 0,
            Dict{Int, Float64}(),
            Dict{Int, UnitRange{Int}}(),
            Dict{Int, Float64}(),
            Dict{Int, UnitRange{Int}}())
    end
end


using StaticLinearMaps
using StaticArrays
using Base.Test

const ℝ{L} = SVector{L,Float64}

# write your own tests here
include("testmain.jl")