module Subspaces

using DocStringExtensions

import Base.hcat, Base.vcat, Base.hvcat, Base.cat, Base.+, Base.*, Base.kron, Base.show, Base.iterate, Base.==, Base.in, Base.adjoint
import Base.|, Base.&, Base.~, Base./, Base.⊆, Base.⊇
import Base.size
using LinearAlgebra
using Convex
using Compat

export Subspace, dim, each_basis_element
export random_subspace, random_hermitian_subspace, empty_subspace, full_subspace
export tobasis, frombasis, random_element, projection, perp
export hermitian_basis
export variable_in_space
export ⟂

default_tol = 1e-6

"""
    Subspace(basis::AbstractArray{AbstractArray{N}, 1}; tol::Real=$(default_tol))

Create a subspace from the given basis, expressed as a list of basis vectors.

    Subspace(basis::AbstractArray{N+1}; tol::Real=$(default_tol))

Create a subspace from a basis given as a multi-dimensional array.  The last array index
enumerates the basis elements.  E.g., if given a matrix this constructor will create a
subspace representing the column span of that matrix.

The `tol` parameter sets the tolerance for determining whether vectors are linearly
dependent.

$(TYPEDFIELDS)

```jldoctest
julia> Subspace([[1, 2, 3], [4, 5, 6]]) == Subspace([ 1 4; 2 5; 3 6])
true
```
"""
struct Subspace{T, N}
    """An orthonormal basis for this subspace.  The final index of this array indexes the
    basis vectors."""
    basis::Array # FIXME can it somehow be Array{T, N+1}?
    """An orthonormal basis for the perpendicular subspace."""
    perp::Array # FIXME can it somehow be Array{T, N+1}?
    """The tolerance for determining whether vectors are linearly dependent."""
    tol::AbstractFloat

    function Subspace(basis::AbstractArray{<:Number}; tol::Real=default_tol)
        shape = size(basis)[1:end-1]
        d = size(basis)[end]
        mat = reshape(basis, prod(shape), d)
        if true
            q = qr(mat, Val(true))
            @compat take = [norm(x) >= tol for x in eachrow(q.R)]
            # This is the most efficient way to get the Q matrix.
            # Multiply by I is slow: https://github.com/JuliaLang/julia/issues/38972
            # Cast to Array changes size: https://github.com/JuliaLang/julia/issues/37102
            qQ = q.Q * Matrix(1.0*I, (size(q.Q)[2], size(q.Q)[2]))
            resize!(take, size(qQ)[2])
            good = qQ[:,take]
            perp = qQ[:,.!take]
        else
            s = svd(mat; full=true)
            lastgood = findlast(s.S .>= tol)
            if typeof(lastgood) == Nothing
                lastgood = 0
            end
            good = s.U[:,1:lastgood]
            perp = s.U[:,lastgood+1:end]
        end
        good = reshape(good, shape..., size(good)[2])
        perp = reshape(perp, shape..., size(perp)[2])
        return new{eltype(good), length(shape)}(good, perp, tol)
    end

    function Subspace(basis::AbstractArray{<:AbstractArray{<:Number}, 1}; tol::Real=default_tol)
        shape = size(basis[1])
        return Subspace(cat(basis...; dims=length(shape)+1); tol=tol)
    end
end

function show(io::IO, S::Subspace)
    print(io, "Subspace{$(eltype(S.basis))} size $(size(S)) dim $(dim(S))")
end

"""
$(TYPEDSIGNATURES)

Returns the size of the elements of a subspace.

```jldoctest
julia> size(Subspace([ [1,2,3], [4,5,6] ]))
(3,)
```
"""
size(S::Subspace)::Tuple{Vararg{<:Integer}} = size(S.basis)[1:end-1]

"""
$(TYPEDSIGNATURES)

Returns the linear dimension of this subspace.

```jldoctest
julia> dim(Subspace([ [1,2,3], [4,5,6] ]))
2
```
"""
dim(S::Subspace)::Integer = size(S.basis)[end]

"""
$(TYPEDSIGNATURES)

Returns the orthogonal subspace.  Can also be written as `~S`.

```jldoctest
julia> S = Subspace([ [1,2,3], [4,5,6] ])
Subspace{Float64} size (3,) dim 2

julia> perp(S)
Subspace{Float64} size (3,) dim 1

julia> perp(S) == ~S
true

julia> perp(S) ⟂ S
true
```
"""
perp(S::Subspace)::Subspace = Subspace(S.perp, tol=S.tol)

"""
$(TYPEDSIGNATURES)

Returns the perpendicular subspace.  Can also be written as `perp(S)`.
"""
(~)(S::Subspace) = perp(S)

# It'd be nice for these to all take SubspaceOrArray but then our overloads seem to be selected even
# when all args are Array.
#const SubspaceOrArray{T, N} = Union{Subspace{T, N}, AbstractArray{T, N}}

"""
$(TYPEDSIGNATURES)

Create a generator that iterates over orthonormal basis vectors of a subspace.
"""
each_basis_element(S::Subspace) = eachslice(S.basis; dims=length(size(S.basis)))

each_basis_element(arr::AbstractArray) = [arr]

function each_basis_element_or_zero(S::Subspace{T, N}) where {T, N}
    if dim(S) == 0
        return [ zeros(T, size(S)) ]
    else
        return each_basis_element(S)
    end
end

"""
    +(a::Subspace, b::Subspace)
    +(a::AbstractArray, b::Subspace)
    +(a::Subspace, b::AbstractArray)
    +(a::UniformScaling, b::Subspace)
    +(a::Subspace, b::UniformScaling)

Linear span of two subspaces, or of a subspace an and array.  Equivalent to |(a, b).
"""
function +(a::Subspace{T, N}, b::Subspace{U, N}) where {T,U,N}
    if size(a) != size(b)
        throw(DimensionMismatch("Array size mismatch: $(size(a)) vs $(size(b))"))
    end
    return Subspace(cat(a.basis, b.basis; dims=N+1), tol=max(a.tol, b.tol))
end

(+)(a::Subspace{T, N}, b::AbstractArray{U, N}) where {T,U,N} = a + Subspace([b], tol=a.tol)
(+)(a::AbstractArray{T, N}, b::Subspace{U, N}) where {T,U,N} = Subspace([a], tol=b.tol) + b
(+)(S::Subspace{T, 2}, x::UniformScaling) where T = S + Subspace([ Array{T}(I, size(S)) ], tol=S.tol)
(+)(x::UniformScaling, S::Subspace{T, 2}) where T = S + I

"""
    |(a::Subspace, b::Subspace)
    |(a::AbstractArray, b::Subspace)
    |(a::Subspace, b::AbstractArray)
    |(a::UniformScaling, b::Subspace)
    |(a::Subspace, b::UniformScaling)

Linear span of two subspaces, or of a subspace an and array.  Equivalent to |(a, b).
"""
(|)(a::Subspace{T, N}, b::Subspace{U, N}) where {T,U,N} = a + b
(|)(a::Subspace{T, N}, b::AbstractArray{U, N}) where {T,U,N} = a + b
(|)(a::AbstractArray{T, N}, b::Subspace{U, N}) where {T,U,N} = a + b
(|)(S::Subspace{T, 2}, x::UniformScaling) where T = S + I
(|)(x::UniformScaling, S::Subspace{T, 2}) where T = S + I

"""
    *(a::Subspace, b::Subspace)
    *(a::AbstractArray, b::Subspace)
    *(a::Subspace, b::AbstractArray)

Linear span of products of elements of space `a` with elements of space `b`.
"""
function *(a::Subspace, b::Subspace)
    #if dim(a) == 0 || dim(b) == 0
    #    return Subspace([ zeros(T, size(a)) * zeros(U, size(b)) ], tol=max(a.tol, b.tol))
    #else
        return Subspace([ x*y for x in each_basis_element_or_zero(a) for y in each_basis_element_or_zero(b) ], tol=max(a.tol, b.tol))
    #end
end

(*)(a::Subspace, b::AbstractArray) = a * Subspace([b], tol=a.tol)
(*)(a::AbstractArray, b::Subspace) = Subspace([a], tol=b.tol) * b

"""
$(TYPEDSIGNATURES)

Linear span of Kronecker products of elements of space `a` with elements of space `b`.
"""
function kron(a::Subspace{T,N}, b::Subspace{U,N}) where {T,U,N}
    return Subspace([
        kron(x, y)
        for x in each_basis_element_or_zero(a)
        for y in each_basis_element_or_zero(b)
    ], tol=max(a.tol, b.tol))
end

kron(a::Subspace, b::AbstractArray) = kron(a, Subspace([b], tol=a.tol))
kron(a::AbstractArray, b::Subspace) = kron(Subspace([a], tol=b.tol), b)

"""
$(TYPEDSIGNATURES)

Vertical concatenations of vector subspaces (direct sum).
"""
vcat(S::Subspace...) = cat(S...; dims=1)

"""
$(TYPEDSIGNATURES)

Horizontal concatenations of vector subspaces (direct sum).
"""
hcat(S::Subspace...) = cat(S...; dims=2)

"""
$(TYPEDSIGNATURES)

Concatenations of vector subspaces (direct sum).
"""
function cat(S::Subspace...; dims)
    n = length(S)
    # FIXME doesn't work well with heterogenous types
    #T = promote_type(map((x)->eltype(x.basis), [S...])...)
    Subspace([
        cat([ i==j ? x : zeros(size(S[i])) for i in 1:n ]...; dims=dims)
        for j in 1:n
        for x in each_basis_element_or_zero(S[j])
    ], tol=maximum([ s.tol for s in S ]))
end

"""
$(TYPEDSIGNATURES)

Concatenations of vector subspaces (direct sum).
"""
function hvcat(rows::Tuple{Vararg{Int}}, S::Subspace{T, N}...) where {T, N}
    n = length(S)
    basis = Array{Array{T, N}, 1}()
    for j in 1:n
        for x in each_basis_element_or_zero(S[j])
            push!(basis, hvcat(rows, [ i==j ? x : zeros(T, size(S[i])) for i in 1:n ]...))
        end
    end
    if isempty(basis)
        push!(basis, hvcat(rows, [ zeros(T, size(S[i])) for i in 1:n ]...))
    end
    return Subspace(basis, tol=maximum([ s.tol for s in S ]))
end

"""
$(TYPEDSIGNATURES)

Adjoint of vector subspace (linear span of adjoints of members of a space).
"""
adjoint(S::Subspace) =
    Subspace([ x' for x in each_basis_element_or_zero(S) ], tol=S.tol)

"""
$(TYPEDSIGNATURES)

Check for membership of a vector in a vector subspace.
"""
function in(x::AbstractArray{<:Number, N}, S::Subspace{<:Number, N}) where N
    return norm(x - projection(S, x)) <= S.tol
end

function in(x::UniformScaling, S::Subspace{T, 2}) where T
    return Matrix{T}(I, size(S)) in S
end

"""
$(TYPEDSIGNATURES)

Check whether `a` is a subspace of `b`.
"""
function in(a::Subspace{<:Number, N}, b::Subspace{<:Number, N}) where N
    shp = size(a)
    if dim(a) > dim(b)
        return false
    end
    Ma = reshape(a.basis, prod(shp), dim(a))
    Mb = reshape(b.basis, prod(shp), dim(b))
    s = svdvals(Mb' * Ma)
    tol = max(a.tol, b.tol)
    return all(s .> (1.0 - tol))
end

"""
$(TYPEDSIGNATURES)

Check whether two vector subspaces are equal.
"""
function ==(a::Subspace{<:Number, N}, b::Subspace{<:Number, N}) where N
    return dim(a) == dim(b) && a in b
end

"""
$(TYPEDSIGNATURES)

Intersection of vector subspaces.
"""
(&)(a::Subspace{T, N}, b::Subspace{U, N}) where {T,U,N} = perp(perp(a) + perp(b))

"""
$(TYPEDSIGNATURES)

Intersection of vector subspaces.
"""
(∩)(a::Subspace{T, N}, b::Subspace{U, N}) where {T,U,N} = a & b

"""
$(TYPEDSIGNATURES)

Quotient of vector subspaces.  Throws error unless `b ⊆ a`.

```jldoctest
julia> a = random_subspace(ComplexF64, 2, 10)
Subspace{Complex{Float64}} size (10,) dim 2

julia> b = random_subspace(ComplexF64, 3, 10)
Subspace{Complex{Float64}} size (10,) dim 3

julia> a / b
ERROR: ArgumentError: divisor must be a subspace of dividend for subspace quotient

julia> (a+b) / a
Subspace{Complex{Float64}} size (10,) dim 3

julia> (a+b) / a == Subspace([ projection(~a, x) for x in each_basis_element(b) ])
true

julia> ((a+b) / b) ⟂ b
true

```
"""
function (/)(a::Subspace{T, N}, b::Subspace{U, N}) where {T,U,N}
    if !(b ⊆ a)
        throw(ArgumentError("divisor must be a subspace of dividend for subspace quotient"))
    end
    return perp(perp(a) + b)
end

(/)(a::Subspace{T, N}, b::AbstractArray{U, N}) where {T,U,N} = a / Subspace([b], tol=a.tol)
(/)(a::Subspace{T, 2}, b::UniformScaling) where T = a / Array{T}(I, size(a))

"""
$(TYPEDSIGNATURES)

Check whether `a` is a subspace of `b`.
"""
(⊆)(a::Subspace{T, N}, b::Subspace{U, N}) where {T,U,N} = a in b

"""
$(TYPEDSIGNATURES)

Check whether `b` is a subspace of `a`.
"""
(⊇)(a::Subspace{T, N}, b::Subspace{U, N}) where {T,U,N} = b in a

"""
$(TYPEDSIGNATURES)

Check whether `a` is orthogonal to `b`.
"""
(⟂)(a::Subspace{T, N}, b::Subspace{U, N}) where {T,U,N} = a ⊆ perp(b)

################
### Math
################

"""
$(TYPEDSIGNATURES)

Returns the basis components of `x` in the basis `S.basis`.

See also: [`frombasis`](@ref).

```jldoctest
julia> S = random_subspace(ComplexF64, 2, 10)
Subspace{Complex{Float64}} size (10,) dim 2

julia> x = randn(10);

julia> size(tobasis(S, x))
(2,)

julia> norm( frombasis(S, tobasis(S, x)) - projection(S, x) ) < 1e-8
true
```
"""
function tobasis(S::Subspace{<:Number, N}, x::AbstractArray{<:Number, N}) where N
    shp = size(S)
    basis_mat = reshape(S.basis, prod(shp), dim(S))
    return basis_mat' * vec(x)
end

"""
$(TYPEDSIGNATURES)

Returns a vector from the given basis components in the basis `S.basis`.

See also: [`tobasis`](@ref).

```jldoctest
julia> S = random_subspace(ComplexF64, 2, 10)
Subspace{Complex{Float64}} size (10,) dim 2

julia> x = randn(2);

julia> size(frombasis(S, x))
(10,)

julia> norm( tobasis(S, frombasis(S, x)) - x ) < 1e-8
true
```
"""
function frombasis(S::Subspace, x::AbstractArray{<:Number, 1})
    shp = size(S)
    basis_mat = reshape(S.basis, prod(shp), dim(S))
    return reshape(basis_mat * x, shp)
end

"""
$(TYPEDSIGNATURES)

Projection of vector `x` onto suspace `S`.
"""
function projection(S::Subspace{<:Number, N}, x::AbstractArray{<:Number, N}) where N
    return frombasis(S, tobasis(S, x))
end

"""
$(TYPEDSIGNATURES)

Returns a random element of a subspace.
"""
random_element(S::Subspace{T}) where T <: Number = frombasis(S, randn(T, dim(S)))

################
### Constructors
################

"""
$(TYPEDSIGNATURES)

Random dimension-`d` subspace of dimension-`siz` vector space, on base field `T`.

```jldoctest
julia> S = random_subspace(ComplexF64, 2, (3, 4))
Subspace{Complex{Float64}} size (3, 4) dim 2
```
"""
function random_subspace(T::Type, d::Int, siz, tol=default_tol)
    if d < 0
        throw(ArgumentError("subspace dimension was negative: $d"))
    elseif d == 0
        return empty_subspace(T, siz, tol)
    else
        b = [ randn(T, siz) for i in 1:d ]
        return Subspace(b, tol=tol)
    end
end

"""
$(TYPEDSIGNATURES)

Random dimension-`d` subspace of `n × n` matrices on base field `T`, satisfying
`x ∈ S => x' ∈ S`.

```jldoctest
julia> S = random_hermitian_subspace(ComplexF64, 2, 3)
Subspace{Complex{Float64}} size (3, 3) dim 2

julia> S == S'
true
```
"""
function random_hermitian_subspace(T::Type, d::Int, n::Int, tol=default_tol)
    if d < 0
        throw(ArgumentError("subspace dimension was negative: $d"))
    elseif d == 0
        return empty_subspace(T, (n, n), tol)
    else
        b = [ randn(T, n, n) for i in 1:d ]
        b = [ x + x' for x in b ]
        return Subspace(b, tol=tol)
    end
end

"""
$(TYPEDSIGNATURES)

Create an empty subspace of dimension-`siz` vector space, on base field `T`.

```jldoctest
julia> S = empty_subspace(ComplexF64, (3, 4))
Subspace{Complex{Float64}} size (3, 4) dim 0
```
"""
empty_subspace(T::Type, siz::Tuple, tol=default_tol) = Subspace([zeros(T, siz)], tol=tol)

"""
$(TYPEDSIGNATURES)

Create an full subspace of dimension-`siz` vector space, on base field `T`.

```jldoctest
julia> S = full_subspace(ComplexF64, (3, 4))
Subspace{Complex{Float64}} size (3, 4) dim 12
```
"""
full_subspace(T::Type, siz::Tuple) = perp(empty_subspace(T, siz))

#############
### Hermitian
#############

function hermit_to_vec(M::AbstractArray{Complex{T}, 2}) where T
    n = size(M)[1]
    size(M)[2] == n || throw(ArgumentError("matrix was not square: $(size(M))"))
    v = zeros(T, n*n)
    k = 0
    for i in 1:n
        v[k+1:k+i] = real(M[1:i, i])
        k += i
        v[k+1:k+i-1] = imag(M[1:i-1, i])
        k += i-1
    end
    return v
end

function vec_to_hermit(v::AbstractArray{T, 1}, n::Integer) where T
    @assert n*n == size(v)[1]
    M = zeros(Complex{T}, n, n)
    k = 0
    for i in 1:n
        M[1:i, i] = v[k+1:k+i]
        M[i, 1:i] = v[k+1:k+i]
        k += i
        M[1:i-1, i] += v[k+1:k+i-1] * 1im
        M[i, 1:i-1] -= v[k+1:k+i-1] * 1im
        k += i-1
    end
    return Hermitian(M)
end

"""
$(TYPEDSIGNATURES)

Given a subspace satisfying `S == S'`, returns a basis in which each basis element is Hermitian.
"""
function hermitian_basis(S::Subspace{Complex{T}})::Array{Hermitian{Complex{T},Array{Complex{T},2}},1} where T
    if dim(S) == 0
        return []
    end
    n = size(S)[1]
    size(S)[2] == n || throw(ArgumentError("subspace size was not square: $(size(S))"))
    M = hcat(
        [ hermit_to_vec( x     + x'     ) for x in each_basis_element(S) ]...,
        [ hermit_to_vec((x*1im)+(x*1im)') for x in each_basis_element(S) ]...
    )
    hb = [ vec_to_hermit(x, n) for x in each_basis_element(Subspace(M)) ]
    @assert (Subspace(hb, tol=S.tol) == S) "Hermitian basis didn't equal original space"
    return hb
end

#########################
### Support for Convex.jl
#########################

"""
$(TYPEDSIGNATURES)

Returns the basis components of `x` in the basis `S.basis`.
"""
function tobasis(S::Subspace{<:Number, N}, x::Convex.AbstractExpr) where N
    shp = size(S)
    basis_mat = reshape(S.basis, prod(shp), dim(S))
    return basis_mat' * vec(x)
end

"""
$(TYPEDSIGNATURES)

Returns a vector from the given basis components in the basis `S.basis`.
"""
function frombasis(S::Subspace{<:Number, 1}, x::Convex.AbstractExpr)
    shp = size(S)
    basis_mat = reshape(S.basis, prod(shp), dim(S))
    return basis_mat * x
end

"""
$(TYPEDSIGNATURES)

Returns a vector from the given basis components in the basis `S.basis`.
"""
function frombasis(S::Subspace{<:Number, 2}, x::Convex.AbstractExpr)
    shp = size(S)
    basis_mat = reshape(S.basis, prod(shp), dim(S))
    return reshape(basis_mat * x, shp...)
end

"""
$(TYPEDSIGNATURES)

Projection of vector `x` onto suspace `S`.
"""
function projection(S::Subspace{<:Number, N}, x::Convex.AbstractExpr) where N
    return frombasis(S, tobasis(S, x))
end

"""
$(TYPEDSIGNATURES)

Constrains the variable `x` to be in the subspace `S`.
"""
function in(x::Convex.AbstractExpr, S::Subspace{<:Number, N}) where N
    return tobasis(perp(S), x) == 0
end

"""
$(TYPEDSIGNATURES)

Creates a `Convex.jl` variable ranging over the given subspace.
"""
function variable_in_space(S::Subspace{<:Complex{<:Real}, N}) where N
    x = Convex.ComplexVariable(dim(S))
    return frombasis(S, x)
end

"""
$(TYPEDSIGNATURES)

Creates a `Convex.jl` variable ranging over the given subspace.
"""
function variable_in_space(S::Subspace{<:Real, N}) where N
    x = Convex.Variable(dim(S))
    return frombasis(S, x)
end

end # module


# ------------------------------------------------------------------------------------------
# # Julia for Data Science - Algorithms
#
# ### Data processing
# In what's next, we will see how to use some of the standard algorithms for data analysis
# implemented in Julia. In particular, we'll look at
#
# 1. Kmeans clustering
# 1. Nearest neighbors with a KDTree
# 1. PCA for dimensionality reduction
# ------------------------------------------------------------------------------------------

using DataFrames, Statistics

# ------------------------------------------------------------------------------------------
# We'll be using the same data for all three of these examples -- the Sacramento real estate
# transactions file that we download next. This is a list of 985 real estate transactions in
# the Sacramento area reported over a five-day period.
# ------------------------------------------------------------------------------------------

download("http://samplecsvs.s3.amazonaws.com/Sacramentorealestatetransactions.csv","houses.csv")
houses = readtable("houses.csv")

# ------------------------------------------------------------------------------------------
# Let's use `Plots` to plot with the `pyplot` backend and start familiarizing ourselves with
# this data set!
# ------------------------------------------------------------------------------------------

using Plots
pyplot()
plot(size=(500,500),leg=false)

# ------------------------------------------------------------------------------------------
# Now let's create a scatter plot to show the price of a house vs. its square footage,
# ------------------------------------------------------------------------------------------

x = houses[!, :sq__ft]
# x = houses[7] # equivalent, useful if file has no header
y = houses[!, :price]
# y = houses[10] # equivalent
scatter(x,y,markersize=3)

# ------------------------------------------------------------------------------------------
# *Houses with 0 square feet that cost money?*
#
# The square footage seems to not have been recorded in these cases.
#
# Filtering these houses out is easy to do!
# ------------------------------------------------------------------------------------------

filter_houses = houses[houses[!, :sq__ft] .> 0, :]  # dot broadcasting
x = filter_houses[!, :sq__ft]
y = filter_houses[!, :price]
scatter(x,y)

# ------------------------------------------------------------------------------------------
# This makes sense! The higher the square footage, the higher the price.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We can filter a `DataFrame` by feature value too, using the `by` function.
# ------------------------------------------------------------------------------------------

by(filter_houses,:type,size)

by(filter_houses,:type,filter_houses->mean(filter_houses[!, :price]))

# ------------------------------------------------------------------------------------------
# ### Example 1: Kmeans Clustering
#
# Now let's do some kmeans clustering on this data.
#
# First, we can load the `Clustering` package to do this.
# ------------------------------------------------------------------------------------------

#Pkg.add("Clustering")
using Clustering

# ------------------------------------------------------------------------------------------
# Let's store the features `:latitude` and `:longitude` in an array `X` that we will pass to
# `kmeans`.
# ------------------------------------------------------------------------------------------

X = filter_houses[!, [:latitude,:longitude]]
X = Array{Float64}(X)

# ------------------------------------------------------------------------------------------
# Each feature is stored as a row of `X`, but we can transpose to make these features
# columns of `X`.
# ------------------------------------------------------------------------------------------

X = transpose(X)

# ------------------------------------------------------------------------------------------
# Now let's plot longitudes vs. latitudes for this data!
# ------------------------------------------------------------------------------------------

location_figure = scatter(X[1, :], X[2, :])
# Alternatively....
# location_figure = scatter(filter_houses[!, :latitude], filter_houses[!, :longitude])
xlabel!("Latitude")
ylabel!("Longitude")
title!("Houses plotted by location")
display(location_figure)

# ------------------------------------------------------------------------------------------
# We want to identify clusters in the data above. As a first pass at guessing how many
# clusters we might need, let's use the number of zip codes in our data.
#
# (Try changing this to see how it impacts results!)
# ------------------------------------------------------------------------------------------

k = length(unique(filter_houses[!, :zip]))

# ------------------------------------------------------------------------------------------
# We can use the `kmeans` function to do kmeans clustering!
# ------------------------------------------------------------------------------------------

C = kmeans(X,k) # try changing k

# ------------------------------------------------------------------------------------------
# Now let's create a new data frame, `df`, with all the same data as `filter_houses` that
# also includes a column for the cluster to which each house has been assigned.
# ------------------------------------------------------------------------------------------

df = DataFrame(cluster = C.assignments,city = filter_houses[!,:city],
    latitude = filter_houses[!,:latitude],longitude = filter_houses[!,:longitude],zip = filter_houses[!,:zip])

# ------------------------------------------------------------------------------------------
# Let's plot each cluster as a different color.
# ------------------------------------------------------------------------------------------

clusters_figure = plot()
for i = 1:k
    # filter df to grab all houses in the ith cluster
    clustered_houses = df[df[!,:cluster].== i,:]
    # grab latitudes and longitudes of all houses in the ith cluster
    xvals = clustered_houses[!,:latitude]
    yvals = clustered_houses[!,:longitude]
    # plot latitudes and longitudes of all houses in the ith cluster
    scatter!(clusters_figure,xvals,yvals,markersize=4)
end
xlabel!("Latitude")
ylabel!("Longitude")
title!("Houses color-coded by cluster")
display(clusters_figure)

# ------------------------------------------------------------------------------------------
# And now let's try coloring them by zip code.
# ------------------------------------------------------------------------------------------

unique_zips = unique(filter_houses[!,:zip])
zips_figure = plot()
for uzip in unique_zips
    # filter houses by zipcode
    subs = filter_houses[filter_houses[!,:zip].==uzip,:]
    # grab the latitudes and longitudes of all houses in a given zipcode/subdivision
    x = subs[!, :latitude]
    y = subs[!, :longitude]
    # plot the houses in this zipcode by latitude and longitude!
    scatter!(zips_figure,x,y)
end
xlabel!("Latitude")
ylabel!("Longitude")
title!("Houses color-coded by zip code")
display(zips_figure)

# ------------------------------------------------------------------------------------------
# Let's see the two plots side by side.
# ------------------------------------------------------------------------------------------

plot(clusters_figure,zips_figure,layout=(2, 1))

# ------------------------------------------------------------------------------------------
# Not exactly! but almost... Now we know that ZIP codes are not randomly assigned!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Example 2: Nearest Neighbor with a KDTree
#
# For this example, let's start by loading the `NearestNeighbors` package.
# ------------------------------------------------------------------------------------------

using NearestNeighbors

# ------------------------------------------------------------------------------------------
# With this package, we'll look for the `knearest` neighbors of one of the houses, `point`.
# ------------------------------------------------------------------------------------------

knearest = 10
id = 70 # try changing this
point = X[:,id]

# ------------------------------------------------------------------------------------------
# Now we can build a `KDTree` and use `knn` to look for `point`'s nearest neighbors!
# ------------------------------------------------------------------------------------------

kdtree = KDTree(X)
idxs, dists = knn(kdtree, point, knearest, true)

# ------------------------------------------------------------------------------------------
# We'll first generate a plot with all of the houses in the same color,
# ------------------------------------------------------------------------------------------

x = filter_houses[!,:latitude];
y = filter_houses[!,:longitude];
scatter(x,y);

# ------------------------------------------------------------------------------------------
# and then overlay the data corresponding to the nearest neighbors of `point` in a different
# color.
# ------------------------------------------------------------------------------------------

x = filter_houses[idxs,:latitude];
y = filter_houses[idxs,:longitude];
scatter!(x,y)

# ------------------------------------------------------------------------------------------
# There are those nearest neighbors in red!
#
# We can see the cities of the neighboring houses by using the indices, `idxs`, and the
# feature, `:city`, to index into the `DataFrame` `filter_houses`.
# ------------------------------------------------------------------------------------------

cities = filter_houses[idxs,:city]

# ------------------------------------------------------------------------------------------
# ### Example 3: PCA for dimensionality reduction
#
# Let us try to reduce the dimensions of the price/area data from the houses dataset.
#
# We can start by grabbing the square footage and prices of the houses and storing them in
# an `Array`.
# ------------------------------------------------------------------------------------------

F = filter_houses[[:sq__ft,:price]]
F = convert(Array{Float64,2},F)'

# ------------------------------------------------------------------------------------------
# Recall how the data looks when we plot housing prices against square footage.
# ------------------------------------------------------------------------------------------

scatter(F[1,:],F[2,:])
xlabel!("Square footage")
ylabel!("Housing prices")

# ------------------------------------------------------------------------------------------
# We can use the `MultivariateStats` package to run PCA
# ------------------------------------------------------------------------------------------

# Pkg.add("MultivariateStats")
using MultivariateStats

# ------------------------------------------------------------------------------------------
# Use `fit` to fit the model
# ------------------------------------------------------------------------------------------

M = fit(PCA, F)

# ------------------------------------------------------------------------------------------
# Note that you can choose the maximum dimension of the new space by setting `maxoutdim`,
# and you can change the method to, for example, `:svd` with the following syntax.
#
# ```julia
# fit(PCA, F; maxoutdim = 1,method=:svd)
# ```
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# It seems like we only get one dimension with PCA! Let's use `transform` to map all of our
# 2D data in `F` to `1D` data with our model, `M`.
# ------------------------------------------------------------------------------------------

y = transform(M, F)

# ------------------------------------------------------------------------------------------
# Let's use `reconstruct` to put our now 1D data, `y`, in a form that we can easily overlay
# (`Xr`) with our 2D data in `F` along the principle direction/component.
# ------------------------------------------------------------------------------------------

Xr = reconstruct(M, y)

# ------------------------------------------------------------------------------------------
# And now we create that overlay, where we can see points along the principle component in
# red.
#
# (Each blue point maps uniquely to some red point!)
# ------------------------------------------------------------------------------------------

scatter(F[1,:],F[2,:])
scatter!(Xr[1,:],Xr[2,:])
xlabel!("Square footage")
ylabel!("Housing prices")


# Form clique-reduction graph
function ExtractTriMat(A::SparseMatrixCSC{Float64,Int64})
    @assert(issymmetric(A))
    n = size(A,1)
    Am = zeros(n,n)
    for i = 1:n-2
        for j = i+1:n-1
            Aji = A[j,i]
            if Aji == 0
                continue
            end

            for k = j+1:n

                if A[k,j]+A[k,i] == 2
                    Am[j,i]+=1/2
                    Am[k,i]+=1/2
                    Am[k,j]+=1/2
                end
            end
        end
    end
    Am = sparse(Am + Am')
    return Am
end

# Form clique-reduction graph
function ExtractAllTriMat(A::SparseMatrixCSC{Float64,Int64})
    n = size(A,1)
    Am = zeros(n,n)
    for i = 1:n-2
        for j = i+1:n-1
            Aji = A[j,i]
            Aij = A[i,j]
            if Aji == 0 || Aij == 0
                continue
            end

            for k = j+1:n

                if A[k,j]+A[k,i]+A[j,k]+A[i,k] == 4
                    Am[j,i]+=1/2
                    Am[k,i]+=1/2
                    Am[k,j]+=1/2
                end
            end
        end
    end
    Am = sparse(Am + Am')
    return Am
end

# Extract triangles as a hyperedge list
function ExtractTriangle_List(A::SparseMatrixCSC{Float64,Int64})
    n = size(A,1)
    H1 = Vector{Int64}()
    H2 = Vector{Int64}()
    H3 = Vector{Int64}()
    for i = 1:n-2
        for j = i+1:n-1
            Aji = A[j,i]
            Aij = A[i,j]
            if Aji == 0 || Aij == 0
                continue
            end

            for k = j+1:n

                if A[k,j]+A[k,i]+A[j,k]+A[i,k] == 4
                    push!(H1,i)
                    push!(H2,j)
                    push!(H3,k)
                end
            end
        end
    end
    H = [H1 H2 H3]
    return H
end

function ExtractFanMat(A::SparseMatrixCSC{Float64,Int64})
    n = size(A,1)
    Am = zeros(n,n)
    for i = 1:n
        for j = i+1:n
            for k = j+1:n
                for l = k+1:n
                    # Yuck, quad loop.
                    ps= [i j k l; i k j l; i l j k; j k i l; j l i k; k l i j]
                    for w = 1:size(ps,1)
                        a = ps[w,1]; b = ps[w,2]; c = ps[w,3]; d = ps[w,4]
                        if A[a,c] > 0 && A[b,c] > 0 && A[a,d] > 0 && A[b,d] > 0 && A[c,a] ==0 && A[d,a] ==0 && A[c,b] ==0 && A[d,b] ==0
                            Am[a,b] += 1/2
                            Am[b,a] += 1/2
                            Am[c,d] += 1/2
                            Am[d,c] += 1/2
                        end
                    end
                end
            end
        end
    end
    Am = sparse(Am+Am')
    return Am
end

function ExtractFanList(A::SparseMatrixCSC{Float64,Int64})
    n = size(A,1)
    H1 = Vector{Int64}()
    H2 = Vector{Int64}()
    H3 = Vector{Int64}()
    H4 = Vector{Int64}()
    for i = 1:n
        for j = i+1:n
            for k = j+1:n
                for l = k+1:n
                    # Yuck, quad loop.
                    ps= [i j k l; i k j l; i l j k; j k i l; j l i k; k l i j]
                    for w = 1:size(ps,1)
                        a = ps[w,1]; b = ps[w,2]; c = ps[w,3]; d = ps[w,4]
                        if A[a,c] > 0 && A[b,c] > 0 && A[a,d] > 0 && A[b,d] > 0 && A[c,a] ==0 && A[d,a] ==0 && A[c,b] ==0 && A[d,b] ==0
                            push!(H1,a)
                            push!(H2,b)
                            push!(H3,c)
                            push!(H4,d)
                        end
                    end
                end
            end
        end
    end
    H = [H1 H2 H3 H4]
    return H
end

function turntounweighted(Aorig::SparseMatrixCSC{Float64,Int64})
    A = copy(Aorig)
    n = size(A,1)
    for i = 1:n
        for j= 1:n
            if A[i,j] > 0
                A[i,j] = 1
                A[j,i] = 1
            end
        end
    end
    return A
end


@testset "eval_strain_on_gauss_pts1" begin
    m = 20
    n = 10
    h = 0.1
    u3 = zeros((m+1)*(n+1))
    a = zeros(m, n)
    b = zeros(m, n)
    for i = 1:m+1
        for j = 1:n+1
            x = (i-1)*h
            y = (j-1)*h
            u3[ i+(j-1)*(m+1) ] = x^2+y^2
            if i<=m && j<=n
                a[i, j] = 2x
                b[i,j] = 2y
            end
        end
    end
    s = eval_strain_on_gauss_pts1(constant(u3), m, n, h)

    S = run(sess, s)
    A = reshape(S[1:4:end,1], m, n)
    B = reshape(S[1:4:end,2], m, n)

    close("all")
    mesh(A)
    mesh(a, color="orange")

    close("all")
    mesh(B)
    mesh(b, color="orange")

end

@testset "compute_fem_stiffness_matrix1" begin
using Revise
using AdFem
using ADCME
    m = 10
    n = 5
    h = 0.1
    U = zeros((m+1)*(n+1))
    for i = 1:m+1
        for j = 1:n+1
            x = (i-1)*h
            y = (j-1)*h
            U[i+(j-1)*(m+1)] = x^2 + y^2
        end
    end

    hmat = zeros(4m*n, 2, 2)
    for i = 1:4*m*n
        hmat[i,:,:] = [2. 1.;3. 4.]
    end
    A = compute_fem_stiffness_matrix1(constant(hmat), m, n, h)
    A2 = compute_fem_stiffness_matrix1([2. 1.;3. 4.], m, n, h)
    sess = Session()
    run(sess, A) - A2

end


function train(X)
  l = size(X,1)
  p = size(X,2)
  W = zeros(l,l)
  for i in 1:p
    W += X[:,i]*X[:,i]'
  end
  W -= p*eye(l,l)
  return W
end

iterHopfield(W::Array{Float64,2}, x::Vector{Float64}) = sign(W*x)

hamming(a,b) = sum(abs(a-b)/2)

function findmatch(Ps::Array{Float64,2}, v::Vector{Float64})
  l = size(Ps,2)
  min = 1
  minidx = -1
  for i in 1:l
    d = hamming(Ps[:,i], v)
    if min == 0 && d == 0
      warn("find match -- double match")
      return -1
    end
    if d < min
      min = d
      minidx = i
    end
  end
  return minidx
end

function searchHopfield(Ps::Array{Float64,2}, v0::Vector{Float64}; W::Array{Float64,2}=zeros(1,0), maxiter::Int=4)
  W2 = Union{}
  if size(W,2) == 0
    W2 = train(Ps)
  else
    W2 = W
  end
  v = deepcopy(v0)
  for i in 1:maxiter
    v1 = iterHopfield(W2,v)
    hamming(v,v1) > 0 ? v = v1 : return findmatch(Ps, v1)
  end
  return -1
end

function basicTestNet()
  N = 20
  a = sign(randn(N)-1)
  b = sign(randn(N)-1)
  c = sign(randn(N))
  Ps = [a';b';c']'
  a1 = deepcopy(a)
  a1[1] = -1*a[1]
  a1[6] = -1*a[6]
  @show searchHopfield(Ps, a1)
  b1 = deepcopy(b)
  b1[6] = -1*b[6]
  @show searchHopfield(Ps, b1)
  return true
end

basicTestNet()


using ProgressMeter # If you didn't install ProgressMeter, do Pkg.add("ProgressMeter")
unirand(a,b) = a + (b-a) * rand(); # define fn returning uniform random number in (a,b)

# make directory to store the positions and the velocities data.
output_dir = "data"
if ! ispath(output_dir); mkdir(output_dir); end

# distane between i th particle and j th particle at periodic boundary condition
function dist(x1::Array{Float64,1} ,x2::Array{Float64,1} ,L::Float64);
    # x[1]: x position of a particle
    # x[2]: y position of a particle
    # L: system size
    dx = abs(x1[1] - x2[1]) % L;
    dx = minimum([dx; L-dx]);
    dy = abs(x1[2] - x2[2]) % L;
    dy = minimum([dy; L-dy]);
    return sqrt(dx^2 + dy^2)
end

# pairwise distance
function ppdist(x::Array{Float64,1} ,y::Array{Float64,1}, num_particle::Int64, L::Float64);
    D = zeros(num_particle, num_particle);
    for i in 1:num_particle-1
            for j in i+1:num_particle
                    D[i,j] = dist([ x[i]; y[i] ], [ x[j]; y[j] ], L);
            end
    end

    return D;
end

# calculate mean angle
function cal_mean_angle(D::Array{Float64,2}, theta::Array{Float64,1}, num_particle::Int64, r::Float64)
    meanangle = zeros(num_particle)
    for iter2 in 1:num_particle
        truth = D[:,iter2] .< r;
        near_particles_angles = theta[truth];
        meanangle[iter2] =
        atan2(
            mean( sin(near_particles_angles) ),
            mean( cos(near_particles_angles) )
        );
    end
    return meanangle;
end

function vicsek()
    # L:             system size
    # num_particle: the number of particles
    # vc:             constant speed
    # dt:             time step
    # steps:         the number of run steps
    # eta:             Order
    # r:             interaction radius
    const L = 5.0; const num_particle = 300; const vc = 1.0;
    const dt = 0.1; const steps = 500; const eta = 0.1; const r = 0.5;

    # initial position of the particles
    # x_i and y_i corespond to the i th positon of particle.
    x = rand(num_particle) * L; y = rand(num_particle) * L;

    # initial angle (-pi, pi) of the particles;
    theta = [unirand(-pi, pi) for i in 1:num_particle];
    #theta = unirand(-pi, pi, num_particle);

    # initial velocity
    vx = vc .* cos(theta); vy = vc .* sin(theta);

    # start fish swimming!
    progress = Progress(steps);
    for iter in 0:steps-1;
        # output configuration
        filename = @sprintf("./data/%05d.csv",iter);
        writecsv(filename, [x y vx vy])

        # pairwise distance
        # D[i,j] is a distance between i th and j th particles.
        D = ppdist(x, y, num_particle, L);
        D += D';

        # calculate mean angle
        meanangle = zeros(num_particle);
        meanangle = cal_mean_angle(D, theta, num_particle, r);

        # update angle with noise
        theta = meanangle .+ 2*pi*eta .* (rand(num_particle) .- 0.5);

        # update positions and velocities;
        vx = vc .* cos(theta); vy = vc .* sin(theta);
        x += vx * dt; y += vy * dt;

        # periodic boundary condition
        #x = (x .+ L) % L; y = (y .+ L) % L;
        x = mod(x,L); y = mod(y,L);

        next!(progress);
    end

end;
@time vicsek();


using Test
using MARY_fRG.Basics

@testset "等边六边型" begin
    hexa1 = EqHexagon(Point2D(0., 0.), 1)
    println(hexa1)
    hexa2 = EqHexagon(Point2D(-2., 3.), 1)
    @test abs(area(hexa2) - 3*sqrt(3) / 2) < 1e-6
end


struct NormalPeakUvD{T} <: UvSpectrumDensity{T}
    # Signal (Normal distribution with area A)
    A::T # Peak area
    UvNormal::Normal{T} # Normal distribution of the peak
    # background (Polynomial of order N)
    bg_poly::Polynomial{2,T}
    bgxl::T # x-value of the left background level
    bgxr::T # x-value of the right background level
    bgll::T # left background level
    bglr::T # right background level
end

function evaluate(d::NormalPeakUvD{T}, x::Real) where {T}
    return d.A * pdf(d.UvNormal, x) + d.bg_poly(x)
end

function NormalPeakUvD(nt::NamedTuple{(:A,:μ,:σ,:bgxl,:bgxr,:bgll,:bglr)})
    T = promote_type(typeof.(values(nt))...)
    poly_coeffs = poly_coeffs_from_levels(T(nt.bgxl), T(nt.bgxr), T(nt.bgll), T(nt.bglr))
    NormalPeakUvD(T(nt.A), Normal(T(nt.μ), T(nt.σ)), Polynomial(poly_coeffs), T(nt.bgxl), T(nt.bgxr), T(nt.bgll), T(nt.bglr))
end

function poly_coeffs_from_levels(bgxl, bgxr, bgll, bglr)
    slope = (bglr - bgll) / (bgxr - bgxl) # linear slope
    mid_offset = bgll + slope * 0.5*(bgxr - bgxl) # offset
    offset = mid_offset - (bgxr + bgxl)/2 * slope
    offset, slope
end

function initial_parameter_guess(d::Type{NormalPeakUvD}, h::StatsBase.Histogram{<:Any, 1}, ::Type{T} = Float64) where {T}
    nh = normalize(h)
    # estimate background parameters through side windows
    bgxl = T(h.edges[1][1])
    bgxr = T(h.edges[1][end])
    Δx_side = (bgxr - bgxl) / 8
    bgll = zero(T)
    bglr = zero(T)
    widths = diff(h.edges[1])
    inv_widths = inv.(widths)
    bgwl = zero(T)
    bgwr = zero(T)
    for i in 1:StatsBase.binindex(h, bgxl + Δx_side)
        bgll += h.weights[i] * inv_widths[i]
        bgwl += widths[i]
    end
    for i in StatsBase.binindex(h, bgxr - Δx_side):length(h.weights)
        bglr += h.weights[i] * inv_widths[i]
        bgwr += widths[i]
    end
    # estimate peak parameters
    μ::T = EmpiricalDistributions._mean(nh)[1]
    σ::T = sqrt(EmpiricalDistributions._var(nh, (μ,))[1])
    A::T = sum(h.weights) - (bgll*bgwl + bglr*bglr)/2
    if A < 0 A = sum(h.weights) end

    parshape = NamedTupleShape(
        A = ScalarShape{T}(),
        μ = ScalarShape{T}(),
        σ = ScalarShape{T}(),
        bgxl = ConstValueShape(bgxl),
        bgxr = ConstValueShape(bgxr),
        bgll = ScalarShape{T}(),
        bglr = ScalarShape{T}()
    )
    p0 = (
        A = A,
        μ = μ,
        σ = σ,
        bgxl = bgxl,
        bgxr = bgxr,
        bgll = bgll,
        bglr = bglr
    )
    lower_bounds = (
        A = T(0),
        μ = μ - σ,
        σ = T(σ/10),
        bgxl = T(bgxl - minimum(widths)),
        bgxr = T(bgxr - minimum(widths)),
        bgll = T(0),
        bglr = T(0)
    )
    upper_bounds = (
        A = T(2*sum(h.weights)),
        μ = μ + σ,
        σ = 3σ,
        bgxl = T(bgxl + minimum(widths)),
        bgxr = T(bgxr + minimum(widths)),
        bgll = T(2bgll),
        bglr = T(2bglr)
    )
    bounds = (
        A = lower_bounds.A..upper_bounds.A,
        μ = lower_bounds.μ..upper_bounds.μ,
        σ = lower_bounds.σ..upper_bounds.σ,
        bgxl = ConstValueDist(bgxl),
        bgxr = ConstValueDist(bgxr),
        bgll = lower_bounds.bgll..upper_bounds.bgll,
        bglr = lower_bounds.bglr..upper_bounds.bglr
    )
    p0, lower_bounds, upper_bounds, parshape, bounds
end

opt_fit(::Type{NormalPeakUvD}, h::Histogram{<:Any, 1}, ::Type{T} = Float64) where {T} =
    opt_fit(RadiationSpectra.NormalPeakUvD, h, initial_parameter_guess(RadiationSpectra.NormalPeakUvD, h, T)[1:4]...)


# sample from a double well in 1D using HMC
using Printf
using Random
using LinearAlgebra
using ForwardDiff
using QuadGK
using BasicMD
using Statistics

include("potentials.jl")

β = 5.0;
x₀ = [-1.0];
seed = 100;
Δt = 1e-1;
n_iters = 10^4; # number of samples

V = x->DoubleWell(x);

# second moment and energy obsrevables
f₁ = x-> x[1]^2;
f₂ = x-> V(x[1]);
observables = (f₁, f₂);

sampler = RWM(V, β, Δt);

Random.seed!(100);
Random.seed!(100);
Xvals,_ = sample_trajectory(x₀, sampler, options=MDOptions(n_iters=n_iters,n_save_iters=n_save_iters));
f₁_estimate = mean(f₁.(Xvals));
f₂_estimate = mean(f₂.(Xvals));
@printf("f₁ estimate with %d samples: %g\n",length(Xvals), f₁_estimate);
@printf("f₂ estimate with %d samples: %g\n", length(Xvals), f₂_estimate);
Random.seed!(100);
observable_samples = sample_observables(x₀, sampler, observables, options=MDOptions(n_iters=n_iters,n_save_iters=n_save_iters));
@printf("f₁ estimate with %d samples: %g\n",length(observable_samples), mean(observable_samples[1,:]));
@printf("f₂ estimate with %d samples: %g\n", length(observable_samples), mean(observable_samples[2,:]));


import Base.mean
export mean

function mean(x::AffineExpr)
  return sum(x) / (x.size[1] * x.size[2])
end

function mean(x::AffineExpr, dim::Int64)
  return sum(x, dim) / x.size[dim]
end


export hsvtorgb

"""
   hsvtorgb(color)

Convert a `color` from HSV space to RGB.
"""
hsvtorgb(color) = begin
    H, S, V = color
    C = V * S
    X = C * (1 - Base.abs((H / 60) % 2 - 1))
    m = V - C
    if 0 ≤ H < 60
        R′, G′, B′ = C, X, 0
    elseif 60 ≤ H < 120
        R′, G′, B′ = X, C, 0
    elseif 120 ≤ H < 180
        R′, G′, B′ = 0, C, X
    elseif 180 ≤ H < 240
        R′, G′, B′ = 0, X, C
    elseif 240 ≤ H < 300
        R′, G′, B′ = X, 0, C
    elseif 300 ≤ H < 360
        R′, G′, B′ = C, 0, X
    else
        R′, G′, B′ = rand(3)
    end
    R, G, B = R′ + m, G′ + m, B′ + m
    [R; G; B]
end


using SMLMData
using NearestNeighbors

"""
    initialdensity = estimatedensities(smld::SMLMData.SMLD2D,
        clusterdata::Vector{Matrix{Float32}}, params::ParamStruct)

Estimate local emitter densities for clusters in `smld` and `clusterdata`.

# Description
The initial local densities `initialdensity` around each pre-cluster present in
`smld`/`clusterdata` are estimated based on the local density of pre-clusters
throughout the entire set of data as well as some of the rate parameters
provided in `params`.
"""
function estimatedensities(smld::SMLMData.SMLD2D,
    clusterdata::Vector{Matrix{Float32}},
    params::ParamStruct)

    # Define some new parameters.
    dutycycle = params.k_on / (params.k_on + params.k_off + params.k_bleach)
    nclusters = length(clusterdata)
    if isempty(smld.nframes)
        maxframe = maximum(smld.framenum)
    else
        maxframe = smld.nframes
    end

    # If only one cluster is present, we should return an answer right away and stop.
    if nclusters == 1
        if size(clusterdata[1], 1) == 1
            initialdensity = 1
        else
            clusterarea = (maximum(clusterdata[1][:, 1]) - minimum(clusterdata[1][:, 1])) *
                          (maximum(clusterdata[1][:, 1]) - minimum(clusterdata[1][:, 1]))
            initialdensity = (1 / clusterarea) *
                             ((params.k_bleach / params.k_off) / (1 - params.p_miss)) /
                             (1 - exp(-params.k_bleach * dutycycle * (maxframe - 1)))
        end
        return initialdensity
    end

    # Determine the center of all clusters assuming each arose from the same
    # emitter.
    clustercenters = zeros(2, nclusters)
    for nn = 1:nclusters
        # Isolate some arrays to improve readability.
        x = clusterdata[nn][:, 1]
        y = clusterdata[nn][:, 2]
        σ_x = clusterdata[nn][:, 3]
        σ_y = clusterdata[nn][:, 4]

        # Compute the cluster centers based on MLE of position.
        clustercenters[1:2, nn] = [sum(x ./ σ_x .^ 2) / sum(1 ./ σ_x .^ 2)
            sum(y ./ σ_y .^ 2) / sum(1 ./ σ_y .^ 2)]
    end

    # Estimate the local cluster density based on the distance to
    # nearest-neighbors.
    kneighbors = minimum([params.nnearestclusters nclusters - 1])
    kdtree = NearestNeighbors.KDTree(clustercenters)
    _, nndist = NearestNeighbors.knn(kdtree, clustercenters, kneighbors + 1)
    nndist = getindex.(nndist, 1:(kneighbors-1))
    clusterdensity = (kneighbors + 1) ./ (pi * nndist .^ 2)

    # Estimate the density of underlying emitters based on cluster density.
    lambda1 = params.k_bleach * dutycycle
    lambda2 = (params.k_on + params.k_off + params.k_bleach) - lambda1
    initialdensity = clusterdensity *
                     (1.0 / dutycycle) * (1.0 / params.k_off) * (1.0 / (1.0 - params.p_miss)) ./
                     ((1.0 / lambda1) * (1.0 - exp(-lambda1 * (maxframe - 1.0))) -
                      (1.0 / lambda2) * (1.0 - exp(-lambda2 * (maxframe - 1.0))))

    return initialdensity
end


using SparseArrays
using CUDA
using Random
using Adapt

function D2_cpu(idata,odata,Nx,Ny,h)
    for i in 1:Nx
        for j in 1:Ny
            if 2 <= i <= Nx-1 && 2 <= j <= Ny - 1
                odata[i,j] = (idata[i-1,j] + idata[i+1,j] + idata[i,j-1] + idata[i,j+1] - 4*idata[i,j])
            end
        end
    end
end

function D2_split_dev(idata,odata,Nx,Ny,h,::Val{TILE_DIM1}, ::Val{TILE_DIM2}) where {TILE_DIM1, TILE_DIM2}
    tidx = threadIdx().x
    tidy = threadIdx().y

    i = (blockIdx().x - 1) * TILE_DIM1 + tidx
    j = (blockIdx().y - 1) * TILE_DIM2 + tidy

    global_index = (i-1)*Ny+j

    if 0 <= i <= Nx && 1 <= j <= Ny
        odata[i,j] = 0
    end

    if 2 <= i <= Nx-1 && 2 <= j <= Ny - 1
        odata[i,j] = (idata[i-1,j] + idata[i+1,j] + idata[i,j-1] + idata[i,j+1] - 4*idata[i,j])
    end

    nothing
end

function matrix_free_cpu(idata,odata,Nx,Ny,h)
    # odata .= 0
    # for i in 1:Nx
    #     for j in 1:Ny
    #         if (i == 1) && (j == 1)
    #             odata[i,j] = (idata[i,j] - 2*idata[i+1,j] + idata[i+2,j] + idata[i,j] - 2*idata[i,j+1] + idata[i,j+1]) / 4
    #         end
    #         if (i == 1) && (j == Ny)
    #             odata[i,j] = (idata[i,j] - 2*idata[i+1,j] + idata[i+2,j] + idata[i,j] - 2*idata[i,j-1] + idata[i,j-1]) / 4
    #         end
    #         if (i == Nx) && (j == 1)
    #             odata[i,j] = (idata[i,j] - 2*idata[i-1,j] + idata[i-2,j] + idata[i,j] - 2*idata[i,j+1] + idata[i,j+1]) / 4
    #         end
    #         # if (i == Nx) && (j == Ny)
    #         #     odata[i,j] = (idata[i,j] - 2*idata[i-1,j] + idata[i-2,j] + idata[i,j] - 2*idata[i,j-1] + idata[i,j-1]) / 4
    #         # end
    #     end
    # end
    odata .= 0
    # alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    alpha1 = alpha2 = -13/h
    alpha3 = alpha4 = -1
    beta = 1
    (i,j) = (1,1)

    odata[i,j] += (idata[i,j] - 2*idata[i+1,j] + idata[i+2,j] + idata[i,j] - 2*idata[i,j+1] + idata[i,j+2]) / 4 # D2

    odata[i,j] += 2 * alpha3 * (( 1.5* idata[i,j] - 2*idata[i+1,j] + 0.5*idata[i+2,j])) / 4 # Neumann

    odata[i,j] += (2 * beta * (1.5 * idata[i,j]) + 2 * alpha1 * (idata[i,j]) * h) / 4 # Dirichlet
    odata[i,j+1] += (2 * beta * (-1 * idata[i,j])) / 2 # Dirichlet
    odata[i,j+2] += (0.5 * beta * (idata[i,j])) / 2 # Dirichlet

    (i,j) = (1,Ny)
    odata[i,j] += (idata[i,j] - 2*idata[i+1,j] + idata[i+2,j] + idata[i,j] - 2*idata[i,j-1] + idata[i,j-2]) / 4 # D2

    odata[i,j] += 2 * alpha3 * (1.5 * idata[i,j] - 2*idata[i+1,j] + 0.5 * idata[i+2,j]) / 4 # Neumann
    odata[i,j] += (2 * beta * (1.5 * idata[i,j]) + 2 * alpha2 * (idata[i,j]) * h) / 4 # Dirichlet
    odata[i,j-1] += (2 * beta * (-1 * idata[i,j])) / 2 # Dirichlet
    odata[i,j-2] += (0.5 * beta * (idata[i,j])) / 2 # Dirichlet

    (i,j) = (Nx,1)
    odata[i,j] += (idata[i,j] - 2*idata[i-1,j] + idata[i-2,j] + idata[i,j] - 2*idata[i,j+1] + idata[i,j+2]) / 4 # D2

    odata[i,j] += 2 * alpha4 * (( 1.5* idata[i,j] - 2*idata[i-1,j] + 0.5*idata[i-2,j])) / 4 # Neumann
    odata[i,j] += (2 * beta * (1.5 * idata[i,j]) + 2 * alpha1 * (idata[i,j]) * h) / 4 # Dirichlet
    odata[i,j+1] += (2 * beta * (-1 * idata[i,j])) / 2 # Dirichlet
    odata[i,j+2] += (0.5 * beta * (idata[i,j])) / 2 # Dirichlet

    (i,j) = (Nx,Ny)
    odata[i,j] += (idata[i,j] - 2*idata[i-1,j] + idata[i-2,j] + idata[i,j] - 2*idata[i,j-1] + idata[i,j-2]) / 4 # D2

    odata[i,j] += 2 * alpha4 * (1.5 * idata[i,j] - 2*idata[i-1,j] + 0.5 * idata[i-2,j]) / 4 # Neumann
    odata[i,j] += (2 * beta * (1.5 * idata[i,j]) + 2 * alpha2 * (idata[i,j]) * h) / 4 # Dirichlet
    odata[i,j-1] += (2 * beta * (-1 * idata[i,j])) / 2 # Dirichlet
    odata[i,j-2] += (0.5 * beta * (idata[i,j])) / 2 # Dirichlet

    (i,j) = (1,2:Ny-1)
    # odata[i,j] .+= ((idata[i,j] .- 2*idata[i+1,j] .+ idata[i+2,j] .+ idata[i,j .- 1] .- 2*idata[i,j] .+ idata[i,j .+ 1] ) ./ 2 # D2
    #                 # .+ 2 .* alpha3 .* (1.5 * idata[i,j] .- 2*idata[i+1,j] .+ 0.5 * idata[i+2,j]) ./ 2
    #                 )
    # odata[i,j] .+= 2 * alpha3 * (1.5 * idata[i,j] .- 2*idata[i+1,j] .+ 0.5 * idata[i+2,j]) ./ 2 # Neumann

    # test @view
    odata[i,j] .+= (view(idata,i,j) .- 2*view(idata,i+1,j) .+ view(idata,i+2,j) .+ view(idata,i,j .-1) .- 2*view(idata,i,j) .+ view(idata,i,j .+ 1)
                    .+ 2 * alpha3 .* (1.5 * view(idata,i,j) .- 2*view(idata,i+1,j) .+ 0.5 * view(idata,i+2,j))) ./ 2

    (i,j) = (Nx,2:Ny-1)
    # odata[i,j] .+= (idata[i,j] .- 2*idata[i-1,j] .+ idata[i-2,j] .+ idata[i,j .- 1] .- 2*idata[i,j] .+ idata[i,j .+ 1]) / 2 # D2
    # odata[i,j] .+= 2 * alpha4 * (1.5 * idata[i,j] .- 2*idata[i-1,j] .+ 0.5 * idata[i-2,j]) ./ 2 # Neumann
    odata[i,j] .+= ((view(idata,i,j) .- 2*view(idata,i-1,j) .+ view(idata,i-2,j) .+ view(idata,i,j.-1) .- 2*view(idata,i,j) .+ view(idata,i,j.+1))
                    .+ 2 * alpha4 * (1.5 * view(idata,i,j) .- 2*view(idata,i-1,j) .+ 0.5 * view(idata,i-2,j))) ./ 2

    (i,j) = (2:Nx-1,1)
    # odata[i,j] .+= (idata[i.-1,j] .- 2*idata[i,j] .+ idata[i.+1,j] .+ idata[i,j] .- 2*idata[i,j+1] .+ idata[i,j+2]) / 2 # D2
    odata[i,j] .+= (view(idata,i .- 1,j) .- 2*view(idata,i,j) .+ view(idata,i .+ 1,j) .+ view(idata,i,j) .- 2*view(idata,i,j+1) .+ view(idata,i,j+2)) ./ 2

    # odata[i,j] .+= (2 * beta * (1.5 * idata[i,j]) .+ 2 * alpha1 * (idata[i,j]) * h) ./ 2 # Dirichlet
    # odata[i,j+1] .+= (2 * beta * (-1 * idata[i,j])) # Dirichlet
    # odata[i,j+2] .+= (0.5 * beta * (idata[i,j]))  # Dirichlet
    odata[i,j] .+= (2 * beta * (1.5 * view(idata,i,j)) .+ 2 * alpha2 * view(idata,i,j) * h) ./ 2
    odata[i,j+1] .+= (2 * beta * (-1 * view(idata,i,j)))
    odata[i,j+2] .+= (0.5 * beta * (view(idata,i,j)))

    (i,j) = (2:Nx-1,Ny)
    # odata[i,j] .+= (idata[i.-1,j] .- 2*idata[i,j] .+ idata[i.+1,j] .+ idata[i,j] .- 2*idata[i,j-1] .+ idata[i,j - 2]) / 2 # D2
    odata[i,j] .+= (view(idata,i .- 1, j) .- 2*view(idata,i,j) .+ view(idata,i .+1,j) .+ view(idata,i,j) .- 2*view(idata,i,j-1) .+ view(idata,i,j-2)) ./ 2

    # odata[i,j] .+= (2 * beta * (1.5 * idata[i,j]) .+ 2 * alpha1 * (idata[i,j]) * h) ./ 2 # Dirichlet
    # # odata[i,j-1] .+= (2 * beta * (-1 * idata[i,j])) # Dirichlet
    # # odata[i,j-2] .+= (0.5 * beta * (idata[i,j])) # Dirichlet
    odata[i,j] .+= (2 * beta * (1.5 * view(idata,i,j)) .+ 2 * alpha1 * view(idata,i,j) * h) ./ 2
    odata[i,j-1] .+= (2 * beta * (-1 * view(idata,i,j)))
    odata[i,j-2] .+= (0.5 * beta * (view(idata,i,j)))
end

function matrix_free_cpu_v2(idata,odata,Nx,Ny,h)
    # odata .= 0
    # for i in 1:Nx
    #     for j in 1:Ny
    #         if (i == 1) && (j == 1)
    #             odata[i,j] = (idata[i,j] - 2*idata[i+1,j] + idata[i+2,j] + idata[i,j] - 2*idata[i,j+1] + idata[i,j+1]) / 4
    #         end
    #         if (i == 1) && (j == Ny)
    #             odata[i,j] = (idata[i,j] - 2*idata[i+1,j] + idata[i+2,j] + idata[i,j] - 2*idata[i,j-1] + idata[i,j-1]) / 4
    #         end
    #         if (i == Nx) && (j == 1)
    #             odata[i,j] = (idata[i,j] - 2*idata[i-1,j] + idata[i-2,j] + idata[i,j] - 2*idata[i,j+1] + idata[i,j+1]) / 4
    #         end
    #         # if (i == Nx) && (j == Ny)
    #         #     odata[i,j] = (idata[i,j] - 2*idata[i-1,j] + idata[i-2,j] + idata[i,j] - 2*idata[i,j-1] + idata[i,j-1]) / 4
    #         # end
    #     end
    # end
    odata .= 0
    # alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    alpha1 = alpha2 = -13/h
    alpha3 = alpha4 = -1
    beta = 1
    (i,j) = (1,1)

    odata[i,j] += (idata[i,j] - 2*idata[i+1,j] + idata[i+2,j] + idata[i,j] - 2*idata[i,j+1] + idata[i,j+2]) / 4 # D2

    odata[i,j] += 2 * alpha3 * (( 1.5* idata[i,j] - 2*idata[i+1,j] + 0.5*idata[i+2,j])) / 4 # Neumann

    odata[i,j] += (2 * beta * (1.5 * idata[i,j]) + 2 * alpha1 * (idata[i,j]) * h) / 4 # Dirichlet
    odata[i,j+1] += (2 * beta * (-1 * idata[i,j])) / 2 # Dirichlet
    odata[i,j+2] += (0.5 * beta * (idata[i,j])) / 2 # Dirichlet

    (i,j) = (1,Ny)
    odata[i,j] += (idata[i,j] - 2*idata[i+1,j] + idata[i+2,j] + idata[i,j] - 2*idata[i,j-1] + idata[i,j-2]) / 4 # D2

    odata[i,j] += 2 * alpha3 * (1.5 * idata[i,j] - 2*idata[i+1,j] + 0.5 * idata[i+2,j]) / 4 # Neumann
    odata[i,j] += (2 * beta * (1.5 * idata[i,j]) + 2 * alpha2 * (idata[i,j]) * h) / 4 # Dirichlet
    odata[i,j-1] += (2 * beta * (-1 * idata[i,j])) / 2 # Dirichlet
    odata[i,j-2] += (0.5 * beta * (idata[i,j])) / 2 # Dirichlet

    (i,j) = (Nx,1)
    odata[i,j] += (idata[i,j] - 2*idata[i-1,j] + idata[i-2,j] + idata[i,j] - 2*idata[i,j+1] + idata[i,j+2]) / 4 # D2

    odata[i,j] += 2 * alpha4 * (( 1.5* idata[i,j] - 2*idata[i-1,j] + 0.5*idata[i-2,j])) / 4 # Neumann
    odata[i,j] += (2 * beta * (1.5 * idata[i,j]) + 2 * alpha1 * (idata[i,j]) * h) / 4 # Dirichlet
    odata[i,j+1] += (2 * beta * (-1 * idata[i,j])) / 2 # Dirichlet
    odata[i,j+2] += (0.5 * beta * (idata[i,j])) / 2 # Dirichlet

    (i,j) = (Nx,Ny)
    odata[i,j] += (idata[i,j] - 2*idata[i-1,j] + idata[i-2,j] + idata[i,j] - 2*idata[i,j-1] + idata[i,j-2]) / 4 # D2

    odata[i,j] += 2 * alpha4 * (1.5 * idata[i,j] - 2*idata[i-1,j] + 0.5 * idata[i-2,j]) / 4 # Neumann
    odata[i,j] += (2 * beta * (1.5 * idata[i,j]) + 2 * alpha2 * (idata[i,j]) * h) / 4 # Dirichlet
    odata[i,j-1] += (2 * beta * (-1 * idata[i,j])) / 2 # Dirichlet
    odata[i,j-2] += (0.5 * beta * (idata[i,j])) / 2 # Dirichlet

    (i,j) = (1,2:Ny-1)
    odata[i,j] .+= (view(idata,i,j) .- 2*view(idata,i+1,j) .+ view(idata,i+2,j) .+ view(idata,i,j .-1) .- 2*view(idata,i,j) .+ view(idata,i,j .+ 1)
                    .+ 2 * alpha3 .* (1.5 * view(idata,i,j) .- 2*view(idata,i+1,j) .+ 0.5 * view(idata,i+2,j))) ./ 2

    (i,j) = (Nx,2:Ny-1)
    odata[i,j] .+= ((view(idata,i,j) .- 2*view(idata,i-1,j) .+ view(idata,i-2,j) .+ view(idata,i,j.-1) .- 2*view(idata,i,j) .+ view(idata,i,j.+1))
                    .+ 2 * alpha4 * (1.5 * view(idata,i,j) .- 2*view(idata,i-1,j) .+ 0.5 * view(idata,i-2,j))) ./ 2

    (i,j) = (2:Nx-1,1)
    odata[i,j] .+= (view(idata,i .- 1,j) .- 2*view(idata,i,j) .+ view(idata,i .+ 1,j) .+ view(idata,i,j) .- 2*view(idata,i,j+1) .+ view(idata,i,j+2)) ./ 2

    odata[i,j] .+= (2 * beta * (1.5 * view(idata,i,j)) .+ 2 * alpha2 * view(idata,i,j) * h) ./ 2
    odata[i,j+1] .+= (2 * beta * (-1 * view(idata,i,j)))
    odata[i,j+2] .+= (0.5 * beta * (view(idata,i,j)))

    (i,j) = (2:Nx-1,Ny)
    odata[i,j] .+= (view(idata,i .- 1, j) .- 2*view(idata,i,j) .+ view(idata,i .+1,j) .+ view(idata,i,j) .- 2*view(idata,i,j-1) .+ view(idata,i,j-2)) ./ 2

    odata[i,j] .+= (2 * beta * (1.5 * view(idata,i,j)) .+ 2 * alpha1 * view(idata,i,j) * h) ./ 2
    odata[i,j-1] .+= (2 * beta * (-1 * view(idata,i,j)))
    odata[i,j-2] .+= (0.5 * beta * (view(idata,i,j)))
end

function matrix_free_cpu_v3(idata,odata,Nx,Ny,h)
    # odata .= 0
    # for i in 1:Nx
    #     for j in 1:Ny
    #         if (i == 1) && (j == 1)
    #             odata[i,j] = (idata[i,j] - 2*idata[i+1,j] + idata[i+2,j] + idata[i,j] - 2*idata[i,j+1] + idata[i,j+1]) / 4
    #         end
    #         if (i == 1) && (j == Ny)
    #             odata[i,j] = (idata[i,j] - 2*idata[i+1,j] + idata[i+2,j] + idata[i,j] - 2*idata[i,j-1] + idata[i,j-1]) / 4
    #         end
    #         if (i == Nx) && (j == 1)
    #             odata[i,j] = (idata[i,j] - 2*idata[i-1,j] + idata[i-2,j] + idata[i,j] - 2*idata[i,j+1] + idata[i,j+1]) / 4
    #         end
    #         # if (i == Nx) && (j == Ny)
    #         #     odata[i,j] = (idata[i,j] - 2*idata[i-1,j] + idata[i-2,j] + idata[i,j] - 2*idata[i,j-1] + idata[i,j-1]) / 4
    #         # end
    #     end
    # end
    odata .= 0
    # alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    alpha1 = alpha2 = -13/h
    alpha3 = alpha4 = -1
    beta = 1
    (i,j) = (1,1)

    odata[i,j] += (idata[i,j] - 2*idata[i+1,j] + idata[i+2,j] + idata[i,j] - 2*idata[i,j+1] + idata[i,j+2]) / 4 # D2

    odata[i,j] += 2 * alpha3 * (( 1.5* idata[i,j] - 2*idata[i+1,j] + 0.5*idata[i+2,j])) / 4 # Neumann

    odata[i,j] += (2 * beta * (1.5 * idata[i,j]) + 2 * alpha1 * (idata[i,j]) * h) / 4 # Dirichlet
    odata[i,j+1] += (2 * beta * (-1 * idata[i,j])) / 2 # Dirichlet
    odata[i,j+2] += (0.5 * beta * (idata[i,j])) / 2 # Dirichlet

    (i,j) = (1,Ny)
    odata[i,j] += (idata[i,j] - 2*idata[i+1,j] + idata[i+2,j] + idata[i,j] - 2*idata[i,j-1] + idata[i,j-2]) / 4 # D2

    odata[i,j] += 2 * alpha3 * (1.5 * idata[i,j] - 2*idata[i+1,j] + 0.5 * idata[i+2,j]) / 4 # Neumann
    odata[i,j] += (2 * beta * (1.5 * idata[i,j]) + 2 * alpha2 * (idata[i,j]) * h) / 4 # Dirichlet
    odata[i,j-1] += (2 * beta * (-1 * idata[i,j])) / 2 # Dirichlet
    odata[i,j-2] += (0.5 * beta * (idata[i,j])) / 2 # Dirichlet

    (i,j) = (Nx,1)
    odata[i,j] += (idata[i,j] - 2*idata[i-1,j] + idata[i-2,j] + idata[i,j] - 2*idata[i,j+1] + idata[i,j+2]) / 4 # D2

    odata[i,j] += 2 * alpha4 * (( 1.5* idata[i,j] - 2*idata[i-1,j] + 0.5*idata[i-2,j])) / 4 # Neumann
    odata[i,j] += (2 * beta * (1.5 * idata[i,j]) + 2 * alpha1 * (idata[i,j]) * h) / 4 # Dirichlet
    odata[i,j+1] += (2 * beta * (-1 * idata[i,j])) / 2 # Dirichlet
    odata[i,j+2] += (0.5 * beta * (idata[i,j])) / 2 # Dirichlet

    (i,j) = (Nx,Ny)
    odata[i,j] += (idata[i,j] - 2*idata[i-1,j] + idata[i-2,j] + idata[i,j] - 2*idata[i,j-1] + idata[i,j-2]) / 4 # D2

    odata[i,j] += 2 * alpha4 * (1.5 * idata[i,j] - 2*idata[i-1,j] + 0.5 * idata[i-2,j]) / 4 # Neumann
    odata[i,j] += (2 * beta * (1.5 * idata[i,j]) + 2 * alpha2 * (idata[i,j]) * h) / 4 # Dirichlet
    odata[i,j-1] += (2 * beta * (-1 * idata[i,j])) / 2 # Dirichlet
    odata[i,j-2] += (0.5 * beta * (idata[i,j])) / 2 # Dirichlet

    # (i,j) = (1,2:Ny-1)
    i = 1
    idata_N = view(idata,1:3,1:Ny)
    # Threads.@threads for j in 2:Ny-1
    @inbounds for j in 2:Ny-1
        odata[1,j] += (idata_N[1,j] - 2*idata_N[2,j] + idata_N[3,j] + idata_N[1,j-1] - 2* idata_N[1,j] + idata_N[1,j+1] + 2 * alpha3 * (1.5 * idata_N[1,j] - 2*idata_N[2,j] + 0.5*idata_N[3,j])) / 2
    end
    # synchronize()

    i = Nx
    idata_S = view(idata,Nx-2:Nx,1:Ny)
    # Threads.@threads for j in 2:Ny-1
    @inbounds for j in 2:Ny-1
        odata[i,j] += (idata_S[3,j] - 2*idata_S[2,j] + idata_S[1,j] + idata_S[3,j-1] - 2* idata_S[3,j] + idata_S[3,j+1] + 2 * alpha4 * (1.5 * idata_S[3,j] - 2*idata_S[2,j] + 0.5*idata_S[1,j])) / 2
    end
    # synchronize()

    j = 1
    idata_W = view(idata,1:Nx,1:3)

    # @inbounds for i in 2:Nx-1
    #     odata[i,j] += (idata_W[i-1,1] - 2*idata_W[i,1] + idata_W[i+1,1] + idata_W[i,1] - 2*idata_W[i,2] + idata_W[i,3]) / 2
    #     # odata[i,j] += (2 * beta * (1.5 * idata_W[i,1]) + 2 * alpha2 * idata_W[i,1] * h) / 2
    #     # odata[i,j+1] += (2 * beta * (-1 * idata_W[i,1]))
    #     # odata[i,j+2] += (0.5 * beta * idata_W[i,1])
    # end

    CPU_W_T = copy(idata_W')
    odata_W_T= zeros(size(CPU_W_T))
    # @inbounds for i in 2:Nx-1
    @inbounds for i in 2:Nx-1
        odata_W_T[1,i] += (CPU_W_T[1,i-1] - 2*CPU_W_T[1,i] + CPU_W_T[1,i+1] + CPU_W_T[1,i] - 2*CPU_W_T[2,i] + CPU_W_T[3,i]) / 2
        odata_W_T[1,i] += (2 * beta * (1.5 * CPU_W_T[1,i]) + 2 * alpha2 * CPU_W_T[1,i] * h) / 2
        odata_W_T[2,i] += (2 * beta * (-1 * CPU_W_T[1,i]))
        odata_W_T[3,i] += (0.5 * beta * CPU_W_T[1,i])
    end

    # odata[:,1] .+= odata_W_T[1,:]

    j = Ny
    idata_E = view(idata,1:Nx,Ny-2:Ny)
    CPU_E_T = copy(idata_E')
    CPU_OUT_E_T = zeros(size(CPU_E_T))
    @inbounds for i in 2:Nx-1
        CPU_OUT_E_T[3,i] += (CPU_E_T[3,i-1] - 2*CPU_E_T[3,i] + CPU_E_T[3,i+1] + CPU_E_T[3,i] - 2*CPU_E_T[2,i] + CPU_E_T[1,i]) / 2
        CPU_OUT_E_T[3,i] += (2 * beta * (1.5 * CPU_E_T[3,i]) + 2 * alpha1 * CPU_E_T[3,i] * h) / 2
        CPU_OUT_E_T[2,i] += (2 * beta * (-1 * CPU_E_T[3,i]))
        CPU_OUT_E_T[1,i] += (0.5 * beta * CPU_E_T[3,i])
    end

    # @inbounds for i in 2:Nx-1
    #     odata[i,j] += (idata_E[i-1,3] - 2*idata_E[i,3] + idata_E[i+1,3] + idata_E[i,3] - 2*idata_E[i,2] + idata_E[i,1]) / 2
    #     odata[i,j] += (2 * beta * (1.5 * idata_E[i,3]) + 2 * alpha1 * idata_E[i,3] * h) / 2
    #     odata[i,j-1] += (2 * beta * (-1 * idata_E[i,3]))
    #     odata[i,j-2] += (0.5 * beta * idata_E[i,3])
    # end

    odata[:,1:3] .+= odata_W_T'
    odata[:,end-2:end] .+= CPU_OUT_E_T'
    nothing
end

function matrix_free_cpu_optimized(GPU_Array,odata,Nx,Ny,h)
    # input as idata_gpu
    # Calculation using CPU

    odata .= 0

    CPU_W = zeros(Nx,3)
    CPU_W_T = zeros(3,Nx)
    CPU_E = zeros(Nx,3)
    CPU_E_T = zeros(3,Nx)
    CPU_N = zeros(3,Ny)
    CPU_S = zeros(3,Ny)

    CPU_OUT_W = zeros(Nx,3)
    CPU_OUT_E = zeros(Nx,3)
    CPU_OUT_N = zeros(3,Ny)
    CPU_OUT_S = zeros(3,Ny)

    copyto!(CPU_W,GPU_Array[:,1:3])
    copyto!(CPU_E,GPU_Array[:,end-2:end])
    copyto!(CPU_N,GPU_Array[1:3,:])
    copyto!(CPU_S,GPU_Array[end-2:end,:])

    # copyto!(CPU_W,view(GPU_Array,1:Nx,1:3))
    # copyto!(CPU_E,view(GPU_Array,1:Nx,Ny-2:Ny))
    # copyto!(CPU_N,view(GPU_Array,1:3,:))
    # copyto!(CPU_S,view(GPU_Array,Nx-2:Nx,:))

    CPU_W_T .= CPU_W'
    CPU_E_T .= CPU_E'
    CPU_OUT_W_T = zeros(3,Ny)
    CPU_OUT_E_T = zeros(3,Ny)

    alpha1 = alpha2 = -13/h
    alpha3 = alpha4 = -1
    beta = 1

    # (i,j) = (1,2:Ny-1)
    i = 1
    # Threads.@threads for j in 2:Ny-1
    # @inbounds for j in 2:Ny-1
    for j in 2:Ny-1
        CPU_OUT_N[1,j] += (CPU_N[1,j] - 2*CPU_N[2,j] + CPU_N[3,j] + CPU_N[1,j-1] - 2* CPU_N[1,j] + CPU_N[1,j+1] + 2 * alpha3 * (1.5 * CPU_N[1,j] - 2*CPU_N[2,j] + 0.5*CPU_N[3,j])) / 2
    end
    synchronize()

    i = Nx
    # # Threads.@threads for j in 2:Ny-1
    # @inbounds for j in 2:Ny-1
    for j in 2:Ny-1
        CPU_OUT_S[3,j] += (CPU_S[3,j] - 2*CPU_S[2,j] + CPU_S[1,j] + CPU_S[3,j-1] - 2* CPU_S[3,j] + CPU_S[3,j+1] + 2 * alpha4 * (1.5 * CPU_S[3,j] - 2*CPU_S[2,j] + 0.5*CPU_S[1,j])) / 2
    end
    # synchronize()

    j = 1

    # @inbounds for i in 2:Nx-1
    for i in 2:Nx-1
        CPU_OUT_W_T[1,i] += (CPU_W_T[1,i-1] - 2*CPU_W_T[1,i] + CPU_W_T[1,i+1] + CPU_W_T[1,i] - 2*CPU_W_T[2,i] + CPU_W_T[3,i]) / 2
        CPU_OUT_W_T[1,i] += (2 * beta * (1.5 * CPU_W_T[1,i]) + 2 * alpha2 * CPU_W_T[1,i] * h) / 2
        CPU_OUT_W_T[2,i] += (2 * beta * (-1 * CPU_W_T[1,i]))
        CPU_OUT_W_T[3,i] += (0.5 * beta * CPU_W_T[1,i])
    end

    # # odata[:,1] .+= odata_W_T[1,:]

    j = Ny
    # @inbounds for i in 2:Nx-1
    for i in 2:Nx-1
        CPU_OUT_E_T[3,i] += (CPU_E_T[3,i-1] - 2*CPU_E_T[3,i] + CPU_E_T[3,i+1] + CPU_E_T[3,i] - 2*CPU_E_T[2,i] + CPU_E_T[1,i]) / 2
        CPU_OUT_E_T[3,i] += (2 * beta * (1.5 * CPU_E_T[3,i]) + 2 * alpha1 * CPU_E_T[3,i] * h) / 2
        CPU_OUT_E_T[2,i] += (2 * beta * (-1 * CPU_E_T[3,i]))
        CPU_OUT_E_T[1,i] += (0.5 * beta * CPU_E_T[3,i])
    end

    (i,j) = (1,1)

    CPU_OUT_N[1,j] += (CPU_W[i,j] - 2*CPU_W[i+1,j] + CPU_W[i+2,j] + CPU_W[i,j] - 2*CPU_W[i,j+1] + CPU_W[i,j+2]) / 4 # D2

    CPU_OUT_N[1,j] += 2 * alpha3 * (( 1.5* CPU_W[i,j] - 2*CPU_W[i+1,j] + 0.5*CPU_W[i+2,j])) / 4 # Neumann

    CPU_OUT_N[1,j] += (2 * beta * (1.5 * CPU_W[i,j]) + 2 * alpha1 * (CPU_W[i,j]) * h) / 4 # Dirichlet
    CPU_OUT_N[1,j+1] += (2 * beta * (-1 * CPU_W[i,j])) / 2 # Dirichlet
    CPU_OUT_N[1,j+2] += (0.5 * beta * (CPU_W[i,j])) / 2 # Dirichlet

    (i,j) = (1,Ny)
    CPU_OUT_N[1,j] += (CPU_E[i,3] - 2*CPU_E[i+1,3] + CPU_E[i+2,3] + CPU_E[i,3] - 2*CPU_E[i,2] + CPU_E[i,1]) / 4 # D2

    CPU_OUT_N[1,j] += 2 * alpha3 * (1.5 * CPU_E[i,3] - 2*CPU_E[i+1,3] + 0.5 * CPU_E[i+2,3]) / 4 # Neumann
    CPU_OUT_N[1,j] += (2 * beta * (1.5 * CPU_E[i,3]) + 2 * alpha2 * (CPU_E[i,3]) * h) / 4 # Dirichlet
    CPU_OUT_N[1,j-1] += (2 * beta * (-1 * CPU_E[i,3])) / 2 # Dirichlet
    CPU_OUT_N[1,j-2] += (0.5 * beta * (CPU_E[i,3])) / 2 # Dirichlet

    (i,j) = (Nx,1)
    CPU_OUT_S[3,j] += (CPU_W[i,j] - 2*CPU_W[i-1,j] + CPU_W[i-2,j] + CPU_W[i,j] - 2*CPU_W[i,j+1] + CPU_W[i,j+2]) / 4 # D2

    CPU_OUT_S[3,j] += 2 * alpha4 * (( 1.5* CPU_W[i,j] - 2*CPU_W[i-1,j] + 0.5*CPU_W[i-2,j])) / 4 # Neumann
    CPU_OUT_S[3,j] += (2 * beta * (1.5 * CPU_W[i,j]) + 2 * alpha1 * (CPU_W[i,j]) * h) / 4 # Dirichlet
    CPU_OUT_S[3,j+1] += (2 * beta * (-1 * CPU_W[i,j])) / 2 # Dirichlet
    CPU_OUT_S[3,j+2] += (0.5 * beta * (CPU_W[i,j])) / 2 # Dirichlet

    (i,j) = (Nx,Ny)
    CPU_OUT_S[3,j] += (CPU_E[Nx,3] - 2*CPU_E[Nx-1,3] + CPU_E[Nx-2,3] + CPU_E[Nx,3] - 2*CPU_E[Nx,2] + CPU_E[Nx,1]) / 4 # D2

    CPU_OUT_S[3,j] += 2 * alpha4 * (1.5 * CPU_E[Nx,3] - 2*CPU_E[Nx-1,3] + 0.5 * CPU_E[Nx-2,3]) / 4 # Neumann
    CPU_OUT_S[3,j] += (2 * beta * (1.5 * CPU_E[Nx,3]) + 2 * alpha2 * (CPU_E[Nx,3]) * h) / 4 # Dirichlet
    CPU_OUT_S[3,j-1] += (2 * beta * (-1 * CPU_E[Nx,3])) / 2 # Dirichlet
    CPU_OUT_S[3,j-2] += (0.5 * beta * (CPU_E[Nx,3])) / 2 # Dirichlet

    # CPU_OUT_cache = CuArray(CPU_OUT_W_T)'
    # CPU_OUT_cache = CuArray(CPU_OUT_E_T)'
    # CuArray(CPU_OUT_W_T')
    # CuArray(CPU_OUT_E_T')
    # @show CPU_OUT_W_T
    # @show CPU_OUT_E_T
    # CPU_OUT_W .= CPU_OUT_W_T'

    # Copy E & W boundary
    copyto!(view(odata,1:Nx,1:3),CuArray(CPU_OUT_W_T)')
    copyto!(view(odata,1:Nx,Ny-2:Ny),CuArray(CPU_OUT_E_T)')

    # Copy N & S boundary
    copyto!(view(odata,1:1,1:Ny),CuArray(CPU_OUT_N[1,:]))
    copyto!(view(odata,Nx:Nx,1:Ny),CuArray(CPU_OUT_S[end,:]))
    nothing
end

function test_matrix_free_boundary(level)
    Nx = Ny = 2^level+1
    h = 1/(Nx-1)
    Random.seed!(0)
    A = randn(Nx,Ny)
    A_sparse = spzeros(Nx,Ny)
    for i in 1:Nx
        for j in 1:Ny
            if !(( 4 <= i <= Nx - 4) && (4 <= j <= Ny-4))
                A_sparse[i,j] = A[i,j]
            end
        end
    end
    odata = spzeros(Nx,Ny)
    odata_v3 = spzeros(Nx,Ny)

    # # make sure matrix_free_cpu_v3 equals matrix_free_cpu
    # odata_test = spzeros(Nx,Ny)
    # matrix_free_cpu_v3(A,odata_test,Nx,Ny,h)
    # matrix_free_cpu(A,odata,Nx,Ny,h)
    # # @show Array(odata)
    # # @show Array(odata_test)
    # @assert odata ≈ odata_test
    # # end test

    # t_cpu = time()
    # iter_times = 1000
    # for i in 1:iter_times
    #     matrix_free_cpu_v2(A_sparse,odata,Nx,Ny,h)
    #     odata .= 0
    #     # matrix_free_cpu_v3(A,odata,Nx,Ny,h)
    # end
    # t_cpu = time() - t_cpu
    # @show t_cpu

    t_cpu_v3 = time()
    iter_times = 1000
    for i in 1:iter_times
        matrix_free_cpu_v3(A,odata_v3,Nx,Ny,h)
        odata_v3 .= 0
        # matrix_free_cpu_v3(A,odata,Nx,Ny,h)
    end
    synchronize()
    t_cpu_v3 = time() - t_cpu_v3
    @show t_cpu_v3

    t_convert = time()
    for i in 1:iter_times
        cu_out = CUDA.CUSPARSE.CuSparseMatrixCSC(odata)
    end
    synchronize()
    t_convert = time() - t_convert
    @show t_convert
    nothing
end

function test_D2_split(level)
    Nx = Ny = 2^level+1
    h = 1/(Nx-1)
    Random.seed!(0)
    A = randn(Nx,Ny)
    A = sparse(A)
    # cu_A = CUDA.CUSPARSE.CuSparseMatrixCSC(A)
    # cu_A = CUDA.CUSPARSE.CuSparseMatrixCSC(A)
    # cu_out = CUDA.CUSPARSE.CuSparseMatrixCSC(spzeros(Nx,Ny))
    cu_A = CuArray(A)
    cu_out = CuArray(zeros(Nx,Ny))
    TILE_DIM_1 = 16
    TILE_DIM_2 = 16
    griddim = (div(Nx,TILE_DIM_1) + 1, div(Ny,TILE_DIM_2) + 1)
	blockdim = (TILE_DIM_1,TILE_DIM_2)

    @cuda threads=blockdim blocks=griddim D2_split_dev(cu_A,cu_out,Nx,Ny,h,Val(TILE_DIM_1), Val(TILE_DIM_2))

    iter_times = 1000

    t_D2_start = time()
    for _ in 1:iter_times
        @cuda threads=blockdim blocks=griddim D2_split_dev(cu_A,cu_out,Nx,Ny,h,Val(TILE_DIM_1), Val(TILE_DIM_2))
    end
    synchronize()
    t_D2 = time() - t_D2_start

    @show t_D2
    nothing
end

function matrix_free_A_v1(idata,odata)
    Nx,Ny = size(idata)
    h = 1/(Nx-1)
    odata_cpu = spzeros(Nx,Ny)
    # odata_gpu = CUDA.CUSPARSE.CuSparseMatrixCSC(spzeros(Nx,Ny))
    odata_gpu = CuArray(zeros(Nx,Ny))
    # idata_gpu = CUDA.CUSPARSE.CuSparseMatrixCSC(idata)
    idata_gpu = CUDA.CuArray(idata)
    TILE_DIM_1 = 16
    TILE_DIM_2 = 16
    griddim = (div(Nx,TILE_DIM_1) + 1, div(Ny,TILE_DIM_2) + 1)
	blockdim = (TILE_DIM_1,TILE_DIM_2)
    @cuda threads=blockdim blocks=griddim D2_split_dev(idata_gpu,odata_gpu,Nx,Ny,h,Val(TILE_DIM_1), Val(TILE_DIM_2))
    matrix_free_cpu_v2(idata,odata_cpu,Nx,Ny,h)
    odata .= adapt(Array,odata_gpu) .+ odata_cpu
end

function matrix_free_A_v2(idata,odata)
    Nx,Ny = size(idata)
    h = 1/(Nx-1)
    # odata_gpu = CuArray(zeros(Nx,Ny))
    # idata_gpu = CUDA.CuArray(idata)
    TILE_DIM_1 = 16
    TILE_DIM_2 = 16
    griddim = (div(Nx,TILE_DIM_1) + 1, div(Ny,TILE_DIM_2) + 1)
	blockdim = (TILE_DIM_1,TILE_DIM_2)
    @cuda threads=blockdim blocks=griddim D2_split_dev(idata,odata,Nx,Ny,h,Val(TILE_DIM_1), Val(TILE_DIM_2))
end

function matrix_free_A_v3(idata,odata,odata_D2_GPU,odata_boundary_GPU)
    # final version for CG
    # odata .= 0
    Nx,Ny = size(idata)
    # odata_D2 = CUDA.zeros(Nx,Ny)
    # odata_boundary = CUDA.zeros(Nx,Ny)
    # odata_boundary = CuArray(zeros(Nx,Ny))
    h = 1/(Nx-1)
    # odata_gpu = CuArray(zeros(Nx,Ny))
    # idata_gpu = CUDA.CuArray(idata)
    TILE_DIM_1 = 16
    TILE_DIM_2 = 16
    griddim = (div(Nx,TILE_DIM_1) + 1, div(Ny,TILE_DIM_2) + 1)
	blockdim = (TILE_DIM_1,TILE_DIM_2)
    @cuda threads=blockdim blocks=griddim D2_split_dev(idata,odata_D2_GPU,Nx,Ny,h,Val(TILE_DIM_1), Val(TILE_DIM_2))
    matrix_free_cpu_optimized(idata,odata_boundary_GPU,Nx,Ny,h)
    synchronize()
    odata .= odata_D2_GPU .+ odata_boundary_GPU
    # copyto!(view(odata,1:3,:), view(odata_D2_GPU,1:3,:) + view(odata_boundary_GPU,1:3,:))
    # return odata
end

function matrix_free_A_v4(idata,odata)
    # odata .= 0
    Nx,Ny = size(idata)
    h = 1/(Nx-1)
    # odata = CUDA.zeros(Nx,Ny)
    TILE_DIM_1 = 16
    TILE_DIM_2 = 16
    griddim = (div(Nx,TILE_DIM_1) + 1, div(Ny,TILE_DIM_2) + 1)
	blockdim = (TILE_DIM_1,TILE_DIM_2)

    CPU_W = zeros(Nx,3)
    CPU_W_T = zeros(3,Nx)
    CPU_E = zeros(Nx,3)
    CPU_E_T = zeros(3,Nx)
    CPU_N = zeros(3,Ny)
    CPU_S = zeros(3,Ny)

    CPU_OUT_W = zeros(Nx,3)
    CPU_OUT_E = zeros(Nx,3)
    CPU_OUT_N = zeros(3,Ny)
    CPU_OUT_S = zeros(3,Ny)

    copyto!(CPU_W,idata[:,1:3])
    copyto!(CPU_E,idata[:,end-2:end])
    copyto!(CPU_N,idata[1:3,:])
    copyto!(CPU_S,idata[end-2:end,:])
    # copyto!(CPU_W,GPU_Array[:,1:3])
    # copyto!(CPU_E,GPU_Array[:,end-2:end])
    # copyto!(CPU_N,GPU_Array[1:3,:])
    # copyto!(CPU_S,GPU_Array[end-2:end,:])

    @cuda threads=blockdim blocks=griddim D2_split_dev(idata,odata,Nx,Ny,h,Val(TILE_DIM_1), Val(TILE_DIM_2))

    CPU_W_T .= CPU_W'
    CPU_E_T .= CPU_E'
    CPU_OUT_W_T = zeros(3,Ny)
    CPU_OUT_E_T = zeros(3,Ny)

    alpha1 = alpha2 = -13/h
    alpha3 = alpha4 = -1
    beta = 1

    i = 1
    # Threads.@threads for j in 2:Ny-1
    # @inbounds for j in 2:Ny-1
    for j in 2:Ny-1
        CPU_OUT_N[1,j] += (CPU_N[1,j] - 2*CPU_N[2,j] + CPU_N[3,j] + CPU_N[1,j-1] - 2* CPU_N[1,j] + CPU_N[1,j+1] + 2 * alpha3 * (1.5 * CPU_N[1,j] - 2*CPU_N[2,j] + 0.5*CPU_N[3,j])) / 2
    end
    synchronize()

    i = Nx
    # # Threads.@threads for j in 2:Ny-1
    # @inbounds for j in 2:Ny-1
    for j in 2:Ny-1
        CPU_OUT_S[3,j] += (CPU_S[3,j] - 2*CPU_S[2,j] + CPU_S[1,j] + CPU_S[3,j-1] - 2* CPU_S[3,j] + CPU_S[3,j+1] + 2 * alpha4 * (1.5 * CPU_S[3,j] - 2*CPU_S[2,j] + 0.5*CPU_S[1,j])) / 2
    end
    # synchronize()

    j = 1

    # @inbounds for i in 2:Nx-1
    for i in 2:Nx-1
        CPU_OUT_W_T[1,i] += (CPU_W_T[1,i-1] - 2*CPU_W_T[1,i] + CPU_W_T[1,i+1] + CPU_W_T[1,i] - 2*CPU_W_T[2,i] + CPU_W_T[3,i]) / 2
        CPU_OUT_W_T[1,i] += (2 * beta * (1.5 * CPU_W_T[1,i]) + 2 * alpha2 * CPU_W_T[1,i] * h) / 2
        CPU_OUT_W_T[2,i] += (2 * beta * (-1 * CPU_W_T[1,i]))
        CPU_OUT_W_T[3,i] += (0.5 * beta * CPU_W_T[1,i])
    end

    # # odata[:,1] .+= odata_W_T[1,:]

    j = Ny
    # @inbounds for i in 2:Nx-1
    for i in 2:Nx-1
        CPU_OUT_E_T[3,i] += (CPU_E_T[3,i-1] - 2*CPU_E_T[3,i] + CPU_E_T[3,i+1] + CPU_E_T[3,i] - 2*CPU_E_T[2,i] + CPU_E_T[1,i]) / 2
        CPU_OUT_E_T[3,i] += (2 * beta * (1.5 * CPU_E_T[3,i]) + 2 * alpha1 * CPU_E_T[3,i] * h) / 2
        CPU_OUT_E_T[2,i] += (2 * beta * (-1 * CPU_E_T[3,i]))
        CPU_OUT_E_T[1,i] += (0.5 * beta * CPU_E_T[3,i])
    end

    (i,j) = (1,1)

    CPU_OUT_N[1,j] += (CPU_W[i,j] - 2*CPU_W[i+1,j] + CPU_W[i+2,j] + CPU_W[i,j] - 2*CPU_W[i,j+1] + CPU_W[i,j+2]) / 4 # D2

    CPU_OUT_N[1,j] += 2 * alpha3 * (( 1.5* CPU_W[i,j] - 2*CPU_W[i+1,j] + 0.5*CPU_W[i+2,j])) / 4 # Neumann

    CPU_OUT_N[1,j] += (2 * beta * (1.5 * CPU_W[i,j]) + 2 * alpha1 * (CPU_W[i,j]) * h) / 4 # Dirichlet
    CPU_OUT_N[1,j+1] += (2 * beta * (-1 * CPU_W[i,j])) / 2 # Dirichlet
    CPU_OUT_N[1,j+2] += (0.5 * beta * (CPU_W[i,j])) / 2 # Dirichlet

    (i,j) = (1,Ny)
    CPU_OUT_N[1,j] += (CPU_E[i,3] - 2*CPU_E[i+1,3] + CPU_E[i+2,3] + CPU_E[i,3] - 2*CPU_E[i,2] + CPU_E[i,1]) / 4 # D2

    CPU_OUT_N[1,j] += 2 * alpha3 * (1.5 * CPU_E[i,3] - 2*CPU_E[i+1,3] + 0.5 * CPU_E[i+2,3]) / 4 # Neumann
    CPU_OUT_N[1,j] += (2 * beta * (1.5 * CPU_E[i,3]) + 2 * alpha2 * (CPU_E[i,3]) * h) / 4 # Dirichlet
    CPU_OUT_N[1,j-1] += (2 * beta * (-1 * CPU_E[i,3])) / 2 # Dirichlet
    CPU_OUT_N[1,j-2] += (0.5 * beta * (CPU_E[i,3])) / 2 # Dirichlet

    (i,j) = (Nx,1)
    CPU_OUT_S[3,j] += (CPU_W[i,j] - 2*CPU_W[i-1,j] + CPU_W[i-2,j] + CPU_W[i,j] - 2*CPU_W[i,j+1] + CPU_W[i,j+2]) / 4 # D2

    CPU_OUT_S[3,j] += 2 * alpha4 * (( 1.5* CPU_W[i,j] - 2*CPU_W[i-1,j] + 0.5*CPU_W[i-2,j])) / 4 # Neumann
    CPU_OUT_S[3,j] += (2 * beta * (1.5 * CPU_W[i,j]) + 2 * alpha1 * (CPU_W[i,j]) * h) / 4 # Dirichlet
    CPU_OUT_S[3,j+1] += (2 * beta * (-1 * CPU_W[i,j])) / 2 # Dirichlet
    CPU_OUT_S[3,j+2] += (0.5 * beta * (CPU_W[i,j])) / 2 # Dirichlet

    (i,j) = (Nx,Ny)
    CPU_OUT_S[3,j] += (CPU_E[Nx,3] - 2*CPU_E[Nx-1,3] + CPU_E[Nx-2,3] + CPU_E[Nx,3] - 2*CPU_E[Nx,2] + CPU_E[Nx,1]) / 4 # D2

    CPU_OUT_S[3,j] += 2 * alpha4 * (1.5 * CPU_E[Nx,3] - 2*CPU_E[Nx-1,3] + 0.5 * CPU_E[Nx-2,3]) / 4 # Neumann
    CPU_OUT_S[3,j] += (2 * beta * (1.5 * CPU_E[Nx,3]) + 2 * alpha2 * (CPU_E[Nx,3]) * h) / 4 # Dirichlet
    CPU_OUT_S[3,j-1] += (2 * beta * (-1 * CPU_E[Nx,3])) / 2 # Dirichlet
    CPU_OUT_S[3,j-2] += (0.5 * beta * (CPU_E[Nx,3])) / 2 # Dirichlet

    synchronize()

    # Copy W & E boundary
    copyto!(view(odata,1:Nx,1:3),view(odata,1:Nx,1:3) + CuArray(CPU_OUT_W_T)')
    copyto!(view(odata,1:Nx,Ny-2:Ny),view(odata,1:Nx,Ny-2:Ny) + CuArray(CPU_OUT_E_T)')

    # view(odata,1,1:Ny) + CuArray(CPU_OUT_N[1,:])
    # view(odata,Nx,1:Ny) + CuArray(CPU_OUT_S[end,:])
    # # @show size(view(odata,1:1,1:Ny))
    # # @show size(CuArray(CPU_OUT_N[1,:]))

    # Copy N & S boundary
    copyto!(view(odata,1,1:Ny),view(odata,1,1:Ny) + CuArray(CPU_OUT_N[1,:]))
    copyto!(view(odata,Nx,1:Ny),view(odata,Nx,1:Ny) + CuArray(CPU_OUT_S[end,:]))
    nothing
end

function CG_GPU_dev(b_reshaped_GPU,x_GPU)
    (Nx,Ny) = size(b_reshaped_GPU)
    odata = CUDA.zeros(Nx,Ny)
    # odata_D2_GPU = CUDA.zeros(Nx,Ny)
    # odata_boundary_GPU = CUDA.zeros(Nx,Ny)
    # matrix_free_A_v3(x_GPU,odata,odata_D2_GPU,odata_boundary_GPU)
    matrix_free_A_v4(x_GPU,odata)
    r_GPU = b_reshaped_GPU - odata
    p_GPU = copy(r_GPU)
    rsold_GPU = sum(r_GPU .* r_GPU)
    Ap_GPU = CUDA.zeros(Nx,Ny)
    # num_steps = 0
    # @show rsold_GPU
    for i in 1:Nx*Ny
    # for i in 1:20
        # @show i
        # @show rsold_GPU
        # matrix_free_A_v3(p_GPU,Ap_GPU,odata_D2_GPU,odata_boundary_GPU)
        # num_steps += 1
        matrix_free_A_v4(p_GPU,Ap_GPU)
        alpha_GPU = rsold_GPU / (sum(p_GPU .* Ap_GPU))
        x_GPU .= x_GPU .+ alpha_GPU * p_GPU
        r_GPU .= r_GPU .- alpha_GPU * Ap_GPU
        rsnew_GPU = sum(r_GPU .* r_GPU)
        if sqrt(rsnew_GPU) <  sqrt(eps(real(eltype(b_reshaped_GPU))))
            break
        end
        p_GPU .= r_GPU .+ (rsnew_GPU/rsold_GPU) * p_GPU
        rsold_GPU = rsnew_GPU
        # @show rsold_GPU
    end
    # @show num_steps
end

function CG_CPU_dev(A,b,x)
    r = b - A * x;
    p = r;
    rsold = r' * r
    # Ap = p
    for i = 1:length(b)
    # for i = 1:2
        # @show i
        # @show rsold
        Ap = A * p;
        alpha = rsold / (p' * Ap)
        x .= x .+ alpha * p;
        r .= r .- alpha * Ap;
        rsnew = r' * r
        if sqrt(rsnew) < sqrt(eps(real(eltype(b))))
              break
        end
        p = r + (rsnew / rsold) * p;
        rsold = rsnew
        # @show rsold
    end
end

function test_CG()
    x = zeros(Nx*Ny)
    x_GPU = CUDA.zeros(Nx,Ny)
    CG_GPU(b_reshaped_GPU,x_GPU)
    CG_CPU(A,b,x)
end

function test_matrix_free_A(level)
    Nx = Ny = 2^level + 1
    h = 1/(Nx-1)
    println("")
    println("Starting Test")
    println("2D Domain Size: $Nx by $Ny")
    Random.seed!(0)
    # idata = sparse(randn(Nx,Ny))
    idata = CuArray(randn(Nx,Ny))
    # odata = spzeros(Nx,Ny)
    odata = CUDA.zeros(Nx,Ny)
    # odata_v4 = CUDA.zeros(Nx,Ny)
    # odata_D2_GPU = CUDA.zeros(Nx,Ny)
    odata_boundary_GPU = CUDA.zeros(Nx,Ny)
    println("Size of the solution matrix (GPU): ", sizeof(idata), " Bytes")

    idata_cpu = zeros(Nx,Ny)
    odata_cpu = spzeros(Nx,Ny)
    odata_D2 = zeros(Nx,Ny)
    copyto!(idata_cpu,idata)
    println("Size of the solution matrix (CPU): ", sizeof(idata_cpu), " Bytes")
    matrix_free_cpu(idata_cpu,odata_cpu,Nx,Ny,h)
    D2_cpu(idata_cpu,odata_D2,Nx,Ny,h)
    odata_cpu + odata_D2

    println("")

    println("Evaluting time to do one A * b")
    println("Timing results in ms")
    println("")

    iter_times = 1000

    #precompile functions

    # odata_boundary = CUDA.zeros(Nx,Ny)

    # matrix_free_A_v3(idata,odata,odata_D2_GPU,odata_boundary_GPU)
    matrix_free_A_v4(idata,odata)
    matrix_free_cpu_v3(idata_cpu,odata_cpu,Nx,Ny,h)
    matrix_free_cpu_optimized(idata,odata_boundary_GPU,Nx,Ny,h)

    @assert odata_cpu ≈ Array(odata_boundary_GPU)

    # Evaluating only D2
    t_start_D2 = time()
    for _ in 1:iter_times
        matrix_free_A_v2(idata,odata)
    end
    synchronize()
    t_D2 = (time() - t_start_D2) * 1000 / iter_times
    @show t_D2
    # End evaluating D2

    # Evaluating only boundary
    t_start_boundary = time()
    for _ in 1:iter_times
        # matrix_free_cpu_v3(idata_cpu,odata_cpu,Nx,Ny,h)
        matrix_free_cpu_v3(idata_cpu,odata_cpu,Nx,Ny,h)
    end
    t_boundary = (time() - t_start_boundary) * 1000 / iter_times
    @show t_boundary
    # End evaluating boundary

    # Evaluating Both in asynchronous way
    t_start_total = time()
    for _ in 1:iter_times
        matrix_free_A_v2(idata,odata)
    end

    for _ in 1:iter_times
        matrix_free_cpu_v3(idata_cpu,odata_cpu,Nx,Ny,h)
    end
    synchronize()
    t_total = (time() - t_start_total) * 1000 / iter_times
    @show t_total
    # End evaluating both in asynchrnous way

    # Test reduction
    t_reduction = time()
    for _ in 1:iter_times
        # reduce(+,idata)
        sum(idata)
        # sum(idata_cpu)
    end
    synchronize()
    t_reduction = (time() - t_reduction ) * 1000 / iter_times
    @show t_reduction
    # End reduction test

    # Evaluating time in Data IO
    t_copy_data = time()
    iter_times_copy_data = 100
    for _ in 1:iter_times_copy_data
        copyto!(idata_cpu,idata)
    end
    # End evaluating time in Data IO
    t_copy_data = ( time() - t_copy_data ) * 1000 / iter_times_copy_data
    @show t_copy_data

    t_boundary_data_part = time()
    iter_times_copy_data = 100
    for _ in 1:iter_times_copy_data
        # matrix_free_cpu_v4(idata,odata_cpu,Nx,Ny,h)
        matrix_free_cpu_optimized(idata,odata_boundary_GPU,Nx,Ny,h)
    end
    # End evaluating time in Data IO
    t_boundary_data_part = ( time() - t_boundary_data_part ) * 1000 / iter_times_copy_data
    @show t_boundary_data_part

     # Evaluating final matrix-free A

    #  t_start_A = time()
    #  for _ in 1:iter_times
    #      matrix_free_A_v3(idata,odata,odata_D2_GPU,odata_boundary_GPU)
    #  end
    #  synchronize()
    #  t_A = (time() - t_start_A) * 1000 / iter_times
    #  @show t_A
     # End evaluating matrix-free A in asynchrnous way

    #  iter_times = 1
    t_start_A_v4 = time()
    for _ in 1:iter_times
         matrix_free_A_v4(idata,odata)
    end
    synchronize()
    t_A_v4 = (time() - t_start_A_v4) * 1000 / iter_times
    @show t_A_v4

    CUDA.unsafe_free!(idata)
    CUDA.unsafe_free!(odata)
    CUDA.unsafe_free!(odata_boundary_GPU)
    nothing
end

# test_matrix_free_A(13)
# test_matrix_free_A(14)
# test_matrix_free_A(15)


using Base.Test
import Roots.fzero

# test problems from Table 1 of paper referenced in fzero.jl
# 1
@test_approx_eq fzero(x -> sin(x) - x/2, pi/2, pi) 1.89549426703398094714
@test_approx_eq Roots.a42a(x -> sin(x) - x/2, pi/2, pi) 1.89549426703398094714

# 2
function test2(x)
    -2*sum([(2i - 5)^2/(x - i^2)^3 for i = 1:20])
end
@test_approx_eq fzero(test2, 1 + 1e-9, 4 - 1e-9) 3.0229153472730568
@test_approx_eq_eps fzero(test2, 100 + 1e-9, 121 - 1e-9) 110.02653274766949 1e-9

# 3
function test3(x, a, b)
    a*x*exp(b*x)
end
@test_approx_eq fzero(x -> test3(x, -40, -1), -9, 31) 0.0
@test_approx_eq fzero(x -> test3(x, -100, -2), -9, 31) 0.0
@test_approx_eq fzero(x -> test3(x, -200, -3), -9, 31) 0.0

# 4
function test4(x, n, a)
    x^n - a
end
@test_approx_eq fzero(x -> test4(x, 4, 0.2), 0, 5) 0.668740304976422
@test_approx_eq fzero(x -> test4(x, 12, 1), 0, 5) 1.0
@test_approx_eq fzero(x -> test4(x, 8, 1), -0.95, 4.05) 1.0
@test_approx_eq fzero(x -> test4(x, 12, 1), -0.95, 4.05) 1.0

# 5
@test_approx_eq fzero(x -> sin(x) - 0.5, 0, 1.5) 0.5235987755982989

# 6
function test6(x, n)
    2x*exp(-n) - 2*exp(-n*x) + 1
end
@test_approx_eq fzero(x -> test6(x, 1), 0, 1) 0.42247770964123665
@test_approx_eq fzero(x -> test6(x, 20), 0, 1) 0.03465735902085385
@test_approx_eq fzero(x -> test6(x, 100), 0, 1) 0.006931471805599453

# 7
function test7(x, n)
    (1 + (1 - n)^2)*x - (1 - n*x)^2
end
@test_approx_eq fzero(x -> test7(x, 5), 0, 1) 0.0384025518406219
@test_approx_eq fzero(x -> test7(x, 10), 0, 1) 0.0099000099980005
@test_approx_eq fzero(x -> test7(x, 20), 0, 1) 0.0024937500390620117

# 8
function test8(x, n)
    x^2 - (1 - x)^n
end
@test_approx_eq fzero(x -> test8(x, 2), 0, 1) 0.5
@test_approx_eq fzero(x -> test8(x, 5), 0, 1) 0.345954815848242
@test_approx_eq fzero(x -> test8(x, 10), 0, 1) 0.24512233375330725
@test_approx_eq fzero(x -> test8(x, 20), 0, 1) 0.16492095727644096

# 9
function test9(x, n)
    (1 + (1 - n)^4)*x - (1 - n*x)^4
end
@test_approx_eq fzero(x -> test9(x, 1), 0, 1) 0.2755080409994844
@test_approx_eq fzero(x -> test9(x, 2), 0, 1) 0.1377540204997422
@test_approx_eq fzero(x -> test9(x, 20), 0, 1) 7.668595122185337e-6

# 10
function test10(x, n)
    exp(-n*x)*(x - 1) + x^n
end
@test_approx_eq fzero(x -> test10(x, 1), 0, 1) 0.401058137541547
@test_approx_eq fzero(x -> test10(x, 5), 0, 1) 0.5161535187579336
@test_approx_eq fzero(x -> test10(x, 20), 0, 1) 0.5527046666784878

# 11
function test11(x, n)
    (n*x - 1)/((n - 1)*x)
end
@test_approx_eq fzero(x -> test11(x, 2), 0.01, 1) 0.5
@test_approx_eq fzero(x -> test11(x, 15), 0.01, 1) 0.0666666666666666
@test_approx_eq fzero(x -> test11(x, 20), 0.01, 1) 0.05

# 12
function test12(x, n)
    x^(1/n) - n^(1/n)
end
@test_approx_eq fzero(x -> test12(x, 2), 1, 100) 2
@test_approx_eq fzero(x -> test12(x, 6), 1, 100) 6
@test_approx_eq fzero(x -> test12(x, 33), 1, 100) 33

# 13
function test13(x)
    if x == 0
        return 0
    else
        return x/exp(x^2)
    end
end
@test_approx_eq fzero(test13, -1, 4) 0

# 14
function test14(x, n)
    if x >= 0
        return n/20*(x/1.5 + sin(x) - 1)
    else
        return -n/20
    end
end
@test_approx_eq fzero(x -> test14(x, 1), -1e4, pi/2) 0.6238065189616123
@test_approx_eq fzero(x -> test14(x, 40), -1e4, pi/2) 0.6238065189616123

# 15
function test15(x, n)
    if x > 2e-3/(1+ n)
        return e - 1.859
    elseif x < 0
        return -0.859
    else
        return exp((n + 1)*x/2*1e3) - 1.859
    end
end
@test_approx_eq fzero(x -> test15(x, 20), -1e4, 1e4) 0.00005905130559421971
@test_approx_eq fzero(x -> test15(x, 40), -1e4, 1e4) 0.000030245790670210097
@test_approx_eq fzero(x -> test15(x, 100), -1e4, 1e4) 0.000012277994232461523
@test_approx_eq fzero(x -> test15(x, 1000), -1e4, 1e4) 1.2388385788997142e-6

## various tests
## issue #29, basically f(a) or f(b) so big we get NaN for guess from bracket.
f = x ->  x + exp(x)
@test_approx_eq fzero(f, 0, [-1e6, 1e6]) -0.5671432904097838

## test infinite range
@test_approx_eq fzero(x -> x, [-Inf, Inf]) 0.0

## issue with fzeros and roots near 'b'
@test 0 <  maximum(fzeros(x -> sin(x) - 1/1000*x, 0, pi)) < pi

## test secant_method_no_bracket
@test_throws Roots.Dithering Roots.secant_method_no_bracket(x -> (cos(x) * cosh(x))^2, 4.5, 5.0)
@test_throws Roots.StateConverged Roots.secant_method_no_bracket(x -> (cos(x) * cosh(x))^2, 4.5, 4.8, ftol=cbrt(eps()))


"""
    tfactors(x::Incremental)

Returns a vector containing the incremental development factors of
the accumulated triangle built out of the incremental triangle `x`.

# Example
```
julia> ti = Incremental([3 5 2; 4 2 0; 1 0 0])
Incremental Triangle {Int64}:
 3  5  2
 4  2  0
 1  0  0

julia> tfactors(ti)
3-element Vector{Float64}:
 2.0
 1.25
 1.0
```
"""
function tfactors(x::Incremental)
    tmp = taccum(x)
    tfactors = ones(Float64, size(x.claims)[2])
    for i in 1:(size(tmp.claims)[2] - 1)
        tfactors[i] = sum(tmp.claims[:,i + 1]) / sum(tmp.claims[1:(end - i), i])
    end
    return tfactors
end

"""
    tfactors(x::Accumulated)

Returns a vector containing the incremental development factors of the accumulated triangle `x`.

# Example
````
julia> ta = Accumulated([3 8 10; 4 6 0; 1 0 0])
Accumulated Triangle {Int64}:
 3  8  10
 4  6   0
 1  0   0

julia> tfactors(ta)
3-element Vector{Float64}:
 2.0
 1.25
 1.0
```
"""
function tfactors(x::Accumulated)
    tfactors = ones(Float64, size(x.claims)[2])
    for i in 1:(size(x.claims)[2] - 1)
        tfactors[i] = sum(x.claims[:,i + 1]) / sum(x.claims[1:(end - i), i])
    end
    return tfactors
end

"""
    fillcl(x::Accumulated)

Returns a completed accumulated triangle, which contains both the
known accumulated claims of accumulated triangle `x` and the chain ladder
estimated accumulated claims.

# Example
```
julia> ta = Accumulated([3 8 10; 4 6 0; 1 0 0])
Accumulated Triangle {Int64}:
 3  8  10
 4  6   0
 1  0   0

julia> fillcl(ta)
Accumulated Completed Triangle {Float64}:
 3.0  8.0  10.0
 4.0  6.0   7.5
 1.0  2.0   2.5
```
"""
function fillcl(x::Accumulated)
    factores = tfactors(x)
    matriz_completa = convert(Matrix{Float64}, deepcopy(x.claims))
    for i in 1:size(matriz_completa, 2)
        for j in 1:size(matriz_completa, 1)
            if (i + j) > (size(matriz_completa, 1) + 1)
                matriz_completa[i, j] = matriz_completa[i, j - 1] * factores[j - 1]
            end
        end
    end
    return AccumulatedCompleted(matriz_completa)
end

"""
    fillcl(x::Incremental)

Returns a completed incremental triangle, which contains both the
known incremental claims of incremental triangle `x` and the chain ladder
estimated incremental claims.

# Example
```
julia> ti = Incremental([3 5 2; 4 2 0; 1 0 0])
Incremental Triangle {Int64}:
 3  5  2
 4  2  0
 1  0  0

julia> fillcl(ti)
Incremental Completed Triangle {Float64}:
 3.0  5.0  2.0
 4.0  2.0  1.5
 1.0  1.0  0.5
```
"""
function fillcl(x::Incremental)
    factores = tfactors(x)
    tri_acum = taccum(x)
    matriz_completa = convert(Matrix{Float64}, deepcopy(tri_acum.claims))
    for i in 1:size(matriz_completa, 2)
        for j in 1:size(matriz_completa, 1)
            if (i + j) > (size(matriz_completa, 1) + 1)
                matriz_completa[i, j] = matriz_completa[i, j - 1] * factores[j - 1]
            end
        end
    end
    des = zeros(Float64, size(matriz_completa, 1), size(matriz_completa, 1))
    des[:,1] = matriz_completa[:,1]
    for j in 2:size(matriz_completa, 1)
        for i in 1:size(matriz_completa, 1)
            des[i, j] = matriz_completa[i, j] - matriz_completa[i, j - 1]
        end
    end
    return IncrementalCompleted(des)
end

"""
    ibnrcl(x::Accumulated)

Returns the estimated incurred but not reported claim reserve for the
accumulated triangle `x` using the chain ladder methodology.

# Example
```
julia> ta = Accumulated([3 8 10; 4 6 0; 1 0 0])
Accumulated Triangle {Int64}:
 3  8  10
 4  6   0
 1  0   0

julia> ibnrcl(ta)
3.0
```
"""
function ibnrcl(x::Accumulated)
    tri_completado = fillcl(x)
    ib = 0
    for i in 1:size(tri_completado.claims, 2)
        j = (tri_completado.claims[i,end] - tri_completado.claims[i, end - i + 1])
        ib = ib + j
    end
    return ib
end

"""
    ibnrcl(x::Incremental)

Returns the estimated incurred but not reported claim reserve for the
incremental triangle `x` using the chain ladder methodology.

# Example
```
julia> ti = Incremental([3 5 2; 4 2 0; 1 0 0])
Incremental Triangle {Int64}:
 3  5  2
 4  2  0
 1  0  0

julia> ibnrcl(ti)
3.0
```
"""
function ibnrcl(x::Incremental)
    y = taccum(x)
    return ibnrcl(y)
end