module Degeneracy

# """
# Degeneracy Module:
#     The local polytope has no preference for input, output, and party labels.
#     As a result, vertices and facet inequalities are computed which are degenerate
#     under permutation of inputs, outputs, or parties. This module houses methods
#     for identifying, enumerating, and eliminating such degeneracies.
# """

using ..Behavior

using LinearAlgebra
using QBase: permutation_matrices

# """
# symmetry_groups:
#     A symmetry group is generated by applying each the transforms onto a "canonical"
#     group element. This function constructs each of the unique groups that can
#     be generated from the set of elements.
#
# Inputs:
#     generators: Array of Matrices, generator of the group symmetry
#     elements: Array of row vectors, "canonical" elements on which a generator acts
#     type: optional String, "facet" or "vertex" denotes what the element is.
#
# Output:
#     symmetry_groups: Array, group elements are stored as arrays of row vectors.
# """
function symmetry_groups(generators, elements; type="facet")

    symmetry_groups = []

    num_generators = length(generators)
    num_transforms = 2^(num_generators)

    transforms = []
    for el in elements

        # check if element exists in symmetry group
        grouped_elements = collect(Iterators.flatten(symmetry_groups))

        if in(el,grouped_elements)
            continue
        end

        # form symmetry groups by applying all transforms to each element
        new_group = []
        for transform in generators

            if type == "facet"
                e = el*transform
            elseif type == "vertex"
                e = transform*el
            end

            if !in(e,new_group)
                push!(new_group, e)
            end
        end

        push!(symmetry_groups, new_group)
    end

    symmetry_groups
end

# """
# merge_symmetry_groups:
#     Combines two sets of groups such that groups with overlapping terms are merged.
#     The resulting sets are the group sets for the combined symmetry generators.
#
# Input:
#     group1/2_sets: Array, contains arrays of group elements. In practice, group1
#         and group2 should contain different partitions of the same elements.
#
# Output:
#     merge_group_sets: Array, contains arrays of group elements which are closed
#         under the generators of group1 and group2.
# """
function merge_symmetry_groups(group1_sets, group2_sets)

    merge_group_sets = group2_sets

    for group1_set in group1_sets

        # get the indices of the group2 sets which overlap with the group1 set
        overlap_ids = []

        merge_group_id = 1
        for merge_group_set in merge_group_sets

            group_overlap_exists = !all(
                (x) -> x === nothing,
                indexin(group1_set, merge_group_set)
            )
            if group_overlap_exists
                push!(overlap_ids, merge_group_id)
            end

            merge_group_id += 1
        end
        # reconstruct group2 s.t. group1[g1_id] is merged with group2[]
        merge_group = unique(cat(group1_set,  collect(Iterators.flatten(merge_group_sets[overlap_ids])), dims=1))
        merge_group_sets = merge_group_sets[ filter( (x) -> !in(x,overlap_ids), 1:end) ]
        push!(merge_group_sets, merge_group)
    end

    merge_group_sets
end

# """
# canonical_permutation_groups:
#     Finds the fundamental groups closed under input/output permutations for the
#     specified bell scenario.
#
# Input:
#     α/β_expt: Tuple, (# inputs, # outputs)
#     facets: Array, the complete set of facets (row vectors) for the bell scenario
#     rep: String, the representation used to write the facet vectors
#
# Output:
#     canonical_groups: Array, contains arrays of facets which represent a group
#         closed under bell permutation transformations.
# """
function canonical_facets(α_expt, β_expt, facets, rep="generalized")

    input_relabels = bipartite_input_relabels(α_expt, β_expt, rep)
    (α_out_relabels, β_out_relabels) = bipartite_output_relabels(α_expt, β_expt, rep)

    α_out_sym_groups = symmetry_groups(α_out_relabels, facets)
    β_out_sym_groups = symmetry_groups(β_out_relabels, facets)

    out_sym_groups = merge_symmetry_groups(α_out_sym_groups, β_out_sym_groups)
    in_sym_groups = symmetry_groups(input_relabels, facets)

    canonical_groups = merge_symmetry_groups(out_sym_groups, in_sym_groups)

    canonical_groups
end

# """
# canonical_vertices(α, β, vertices, rep):
#
#     Uses the permutation symmetries to find the vertex-transitive groups.
#
# Input:
#     α/β: Tuple (# inputs, # outputs)
#     vertices: Array, contains a complete set of polytope vertices
#     rep: String, specifies the representation of the vertices
#
# Output:
#     canonical_groups: Array, contains vertices grouped by vertex transitivity
#                       with respect to the provided permutation symmetry groups.
# """
function canonical_vertices(α, β, vertices, rep="generalized")

    input_relabels = bipartite_input_relabels(α, β, rep)
    (α_out_relabels, β_out_relabels) = bipartite_output_relabels(α, β, rep)

    α_out_sym_groups = symmetry_groups(α_out_relabels, vertices, type="vertex")
    β_out_sym_groups = symmetry_groups(β_out_relabels, vertices, type="vertex")

    out_sym_groups = merge_symmetry_groups(α_out_sym_groups, β_out_sym_groups)
    in_sym_groups = symmetry_groups(input_relabels, vertices, type="vertex")

    canonical_groups = merge_symmetry_groups(out_sym_groups, in_sym_groups)

    canonical_groups
end

# """
# bipartite_input_relabels:
#   Returns the set of input relabeling matrices for the specified bipartite
#   bell experiment. A relabeling matrix switches the index of an input such that
#   (x <-> x') and/or (y <-> y').
#
# Inputs:
#   (α/β)_expt: Tuple, experiment parameters (Int::num_inputs, Int::num_outputs)
#   rep: String, specify the matrix representation used
#
# Output:
#   maps: Array, contains the complete set of matrices which transform a behavior
#         to its relabeled counterpart (x <-> x') and/or (y <-> y')
# """
function bipartite_input_relabels(α_expt, β_expt, rep="generalized")
    (α_num_inputs, α_num_outputs) = α_expt
    (β_num_inputs, β_num_outputs) = β_expt

    gen_maps = generalized_input_relabels(α_expt, β_expt, true)

    maps = []

    if rep == "generalized"
        maps = gen_maps
    else
        proj = []
        proj_inv = []

        if rep == "non-signaling"
            proj = Behavior.gen_to_ns_proj(α_expt,β_expt)
            proj_inv = Behavior.ns_to_gen_proj(α_expt,β_expt)
        elseif rep == "fixed-direction"
            proj = Behavior.gen_to_fd_proj(α_expt,β_expt)
            proj_inv = Behavior.fd_to_gen_proj(α_expt,β_expt)
        elseif rep == "normalized"
            proj = Behavior.gen_to_norm_proj(α_expt,β_expt)
            proj_inv = Behavior.norm_to_gen_proj(α_expt,β_expt)
        end

        for map in gen_maps
            push!(maps, proj*map*proj_inv)
        end
    end

    maps
end

# """
# generalized_input_relabels:
#     Constructs relabeling maps for the generalized subspace. Maps are constructed
#     by enumerating all permutations of inputs.
#
# Inputs:
#   expt: Tuple, experiment paramterization (num_inputs, num_outputs)
#   include_constant: Bool, regularization constant is added to relabeling maps
#
# Output:
#   maps: Array, contains relabeling matrices for the generalized subspace
# """
function generalized_input_relabels(α_expt, β_expt, include_constant=false)
    (α_num_inputs, α_num_outputs) = α_expt
    (β_num_inputs, β_num_outputs) = β_expt

    α_maps = permutation_matrices(α_num_inputs)
    β_maps = permutation_matrices(β_num_inputs)

    α_id = diagm(0 => fill(1,α_num_outputs))
    β_id = diagm(0 => fill(1,β_num_outputs))

    maps = []
    for α_map in α_maps
        for β_map in β_maps

            map = kron(α_id, kron(β_id, kron(α_map,β_map)))
            if include_constant
                map = cat([[1][:,:], map]..., dims=(1,2))
            end

            push!(maps, map)
        end
    end

    maps
end

# """
# bipartite_output_relabels(α_expt, β_expt, rep)
#     Enumerates the generalized output relabeling maps and projects the set onto
#     the specified subspace representation.
#
# Inputs:
#     (α/β)_expt: Tuple, the relabeling map
#     rep: String, subspace representation
#
# Outputs:
#     relabels: Array, set of unique output relabeling maps
# """
function bipartite_output_relabels(α_expt, β_expt, rep="normalized")

    (α_num_inputs, α_num_outputs) = α_expt
    (β_num_inputs, β_num_outputs) = β_expt

    α_relabels = []
    β_relabels = []

    gen_relabels = generalized_output_relabels(α_expt,β_expt,true)

    if rep == "generalized"
        (α_relabels, β_relabels) = gen_relabels
    else
        proj = []
        proj_inv = []

        if rep == "non-signaling"
            proj = Behavior.gen_to_ns_proj(α_expt,β_expt)
            proj_inv = Behavior.ns_to_gen_proj(α_expt,β_expt)
        elseif rep == "fixed-direction"
            proj = Behavior.gen_to_fd_proj(α_expt,β_expt)
            proj_inv = Behavior.fd_to_gen_proj(α_expt,β_expt)
        elseif rep == "normalized"
            proj = Behavior.gen_to_norm_proj(α_expt,β_expt)
            proj_inv = Behavior.norm_to_gen_proj(α_expt,β_expt)
        end

        for α_relabel in gen_relabels[1]
            push!(α_relabels, proj*α_relabel*proj_inv)
        end

        for β_relabel in gen_relabels[2]
            push!(β_relabels, proj*β_relabel*proj_inv)
        end
    end

    relabels = (unique(α_relabels), unique(β_relabels))

    relabels
end

# """
# generalized_output_relabels:
#     enumerates the output relabeling maps in the generalized representation.
#
# Inputs:
#     α/β_expt: Tuple, experiment parameters
#     include_constant: Boolean, include regularization term
#
# Outputs:
#   α_maps: Array of output relabeling matrices for alice's outputs
#   β_maps: Array of output relabeling matrices for bob's outputs
# """
function generalized_output_relabels(α_expt,β_expt, include_constant)
    (α_num_inputs, α_num_outputs) = α_expt
    (β_num_inputs, β_num_outputs) = β_expt

    α_perm_maps = permutation_matrices(α_num_outputs)
    β_perm_maps = permutation_matrices(β_num_outputs)

    α_in_id = diagm( 0 => fill(1,α_num_inputs))
    β_in_id = diagm( 0 => fill(1,β_num_inputs))
    α_out_id = diagm( 0 => fill(1,α_num_outputs))
    β_out_id = diagm( 0 => fill(1,β_num_outputs))

    α_maps = []
    β_maps = []
    num_inputs = α_num_inputs*β_num_inputs

    # conditional output relabelings for alice
    α_map_index = []
    for x in 1:α_num_inputs

        x_m = zeros(Int64,(α_num_inputs,α_num_inputs))
        x_m[x,x] = 1

        x_rest = α_in_id - x_m
        αx_maps = []
        for α_perm in α_perm_maps

            map = kron(α_perm, kron(β_out_id,kron(x_m,β_in_id))) + kron(α_out_id, kron(β_out_id,kron(x_rest,β_in_id)))
            if include_constant
                map = cat([[1][:,:], map]..., dims=(1,2))
            end

            push!(αx_maps, map)
        end

        push!(α_map_index, αx_maps)
    end

    if α_num_outputs == 1
        push!(α_maps, α_map_index[1][1])
    else
        α_num_transforms = length(α_perm_maps)^α_num_inputs
        for i in 0:(α_num_transforms - 1)
            id_array = digits(i, base = length(α_perm_maps), pad = α_num_inputs) .+ 1
            push!(α_maps, prod(map(
                (id) -> begin
                    (α_map_index[id])[id_array[id]]
                end,
                1:α_num_inputs
            )))
        end
    end
    # conditional output relabelings for bob

    β_map_index = []
    for y in 1:β_num_inputs

        y_m = zeros(Int64,(β_num_inputs,β_num_inputs))
        y_m[y,y] = 1

        y_rest = β_in_id - y_m
        βy_maps = []
        for β_perm in β_perm_maps

            map = kron(α_out_id, kron(β_perm,kron(α_in_id,y_m))) + kron(α_out_id, kron(β_out_id,kron(α_in_id,y_rest)))
            if include_constant
                map = cat([[1][:,:], map]..., dims=(1,2))
            end

            push!(βy_maps, map)
        end

        push!(β_map_index, βy_maps)
    end

    β_num_transforms = length(β_perm_maps)^β_num_inputs
    if β_num_outputs == 1
        push!(β_maps, β_map_index[1][1])
    else
        for i in 0:(β_num_transforms - 1)
            id_array = digits(i, base = length(β_perm_maps), pad = β_num_inputs) .+ 1
            push!(β_maps, prod(map(
                (id) -> begin
                    (β_map_index[id])[id_array[id]]
                end,
                1:β_num_inputs
            )))
        end
    end

    (α_maps, β_maps)
end

end


# Ordinary metrics

###########################################################
#
#   Metric types
#
###########################################################

type Euclidean <: Metric end
type SqEuclidean <: SemiMetric end
type Chebyshev <: Metric end
type Cityblock <: Metric end
type Jaccard <: Metric end
type RogersTanimoto <: Metric end

immutable Minkowski{T <: Real} <: Metric
    p::T
end

type Hamming <: Metric end

type CosineDist <: SemiMetric end
type CorrDist <: SemiMetric end

type ChiSqDist <: SemiMetric end
type KLDivergence <: PreMetric end
type JSDivergence <: SemiMetric end

type SpanNormDist <: SemiMetric end

typealias UnionMetrics Union{Euclidean, SqEuclidean, Chebyshev, Cityblock, Minkowski, Hamming, Jaccard, RogersTanimoto, CosineDist, CorrDist, ChiSqDist, KLDivergence, JSDivergence, SpanNormDist}

###########################################################
#
#  Define Evaluate
#
###########################################################

function evaluate(d::UnionMetrics, a::AbstractArray, b::AbstractArray)
    if length(a) != length(b)
        throw(DimensionMismatch("first array has length $(length(a)) which does not match the length of the second, $(length(b))."))
    end
    if length(a) == 0
        return zero(result_type(d, a, b))
    end
    s = eval_start(d, a, b)
    if size(a) == size(b)
        @simd for I in eachindex(a, b)
            @inbounds ai = a[I]
            @inbounds bi = b[I]
            s = eval_reduce(d, s, eval_op(d, ai, bi))
        end
    else
        for (Ia, Ib) in zip(eachindex(a), eachindex(b))
            @inbounds ai = a[Ia]
            @inbounds bi = b[Ib]
            s = eval_reduce(d, s, eval_op(d, ai, bi))
        end
    end
    return eval_end(d, s)
end
result_type{T1, T2}(dist::UnionMetrics, ::AbstractArray{T1}, ::AbstractArray{T2}) =
    typeof(eval_end(dist, eval_op(dist, one(T1), one(T2))))
eval_start(d::UnionMetrics, a::AbstractArray, b::AbstractArray) =
    zero(result_type(d, a, b))
eval_end(d::UnionMetrics, s) = s

evaluate{T <: Number}(dist::UnionMetrics, a::T, b::T) = eval_end(dist, eval_op(dist, a, b))

# SqEuclidean
@inline eval_op(::SqEuclidean, ai, bi) = abs2(ai - bi)
@inline eval_reduce(::SqEuclidean, s1, s2) = s1 + s2

sqeuclidean(a::AbstractArray, b::AbstractArray) = evaluate(SqEuclidean(), a, b)
sqeuclidean{T <: Number}(a::T, b::T) = evaluate(SqEuclidean(), a, b)

# Euclidean
@inline eval_op(::Euclidean, ai, bi) = abs2(ai - bi)
@inline eval_reduce(::Euclidean, s1, s2) = s1 + s2
eval_end(::Euclidean, s) = sqrt(s)
euclidean(a::AbstractArray, b::AbstractArray) = evaluate(Euclidean(), a, b)
euclidean(a::Number, b::Number) = evaluate(Euclidean(), a, b)

# Cityblock
@inline eval_op(::Cityblock, ai, bi) = abs(ai - bi)
@inline eval_reduce(::Cityblock, s1, s2) = s1 + s2
cityblock(a::AbstractArray, b::AbstractArray) = evaluate(Cityblock(), a, b)
cityblock{T <: Number}(a::T, b::T) = evaluate(Cityblock(), a, b)

# Chebyshev
@inline eval_op(::Chebyshev, ai, bi) = abs(ai - bi)
@inline eval_reduce(::Chebyshev, s1, s2) = max(s1, s2)
# if only NaN, will output NaN
@inline eval_start(::Chebyshev, a::AbstractArray, b::AbstractArray) = abs(a[1] - b[1])
chebyshev(a::AbstractArray, b::AbstractArray) = evaluate(Chebyshev(), a, b)
chebyshev{T <: Number}(a::T, b::T) = evaluate(Chebyshev(), a, b)

# Minkowski
@inline eval_op(dist::Minkowski, ai, bi) = abs(ai - bi) .^ dist.p
@inline eval_reduce(::Minkowski, s1, s2) = s1 + s2
eval_end(dist::Minkowski, s) = s .^ (1/dist.p)
minkowski(a::AbstractArray, b::AbstractArray, p::Real) = evaluate(Minkowski(p), a, b)
minkowski{T <: Number}(a::T, b::T, p::Real) = evaluate(Minkowski(p), a, b)

# Hamming
@inline eval_op(::Hamming, ai, bi) = ai != bi ? 1 : 0
@inline eval_reduce(::Hamming, s1, s2) = s1 + s2
hamming(a::AbstractArray, b::AbstractArray) = evaluate(Hamming(), a, b)
hamming{T <: Number}(a::T, b::T) = evaluate(Hamming(), a, b)

# Cosine dist
function eval_start{T<:AbstractFloat}(::CosineDist, a::AbstractArray{T}, b::AbstractArray{T})
    zero(T), zero(T), zero(T)
end
@inline eval_op(::CosineDist, ai, bi) = ai * bi, ai * ai, bi * bi
@inline function eval_reduce(::CosineDist, s1, s2)
    a1, b1, c1 = s1
    a2, b2, c2 = s2
    return a1 + a2, b1 + b2, c1 + c2
end
function eval_end(::CosineDist, s)
    ab, a2, b2 = s
    max(1 - ab / (sqrt(a2) * sqrt(b2)), zero(eltype(ab)))
end
cosine_dist(a::AbstractArray, b::AbstractArray) = evaluate(CosineDist(), a, b)

# Correlation Dist
_centralize(x::AbstractArray) = x .- mean(x)
evaluate(::CorrDist, a::AbstractArray, b::AbstractArray) = cosine_dist(_centralize(a), _centralize(b))
corr_dist(a::AbstractArray, b::AbstractArray) = evaluate(CorrDist(), a, b)
result_type(::CorrDist, a::AbstractArray, b::AbstractArray) = result_type(CosineDist(), a, b)

# ChiSqDist
@inline eval_op(::ChiSqDist, ai, bi) = abs2(ai - bi) / (ai + bi)
@inline eval_reduce(::ChiSqDist, s1, s2) = s1 + s2
chisq_dist(a::AbstractArray, b::AbstractArray) = evaluate(ChiSqDist(), a, b)

# KLDivergence
@inline eval_op(::KLDivergence, ai, bi) = ai > 0 ? ai * log(ai / bi) : zero(ai)
@inline eval_reduce(::KLDivergence, s1, s2) = s1 + s2
kl_divergence(a::AbstractArray, b::AbstractArray) = evaluate(KLDivergence(), a, b)

# JSDivergence
@inline function eval_op{T}(::JSDivergence, ai::T, bi::T)
    u = (ai + bi) / 2
    ta = ai > 0 ? ai * log(ai) / 2 : zero(log(one(T)))
    tb = bi > 0 ? bi * log(bi) / 2 : zero(log(one(T)))
    tu = u > 0 ? u * log(u) : zero(log(one(T)))
    ta + tb - tu
end
@inline eval_reduce(::JSDivergence, s1, s2) = s1 + s2
js_divergence(a::AbstractArray, b::AbstractArray) = evaluate(JSDivergence(), a, b)

# SpanNormDist
function eval_start(::SpanNormDist, a::AbstractArray, b::AbstractArray)
    a[1] - b[1], a[1] - b[1]
end
@inline eval_op(::SpanNormDist, ai, bi)  = ai - bi
@inline function eval_reduce(::SpanNormDist, s1, s2)
    min_d, max_d = s1
    if s2 > max_d
        max_d = s2
    elseif s2 < min_d
        min_d = s2
    end
    return min_d, max_d
end

eval_end(::SpanNormDist, s) = s[2] - s[1]
spannorm_dist(a::AbstractArray, b::AbstractArray) = evaluate(SpanNormDist(), a, b)
function result_type{T1, T2}(dist::SpanNormDist, ::AbstractArray{T1}, ::AbstractArray{T2})
    typeof(eval_op(dist, one(T1), one(T2)))
end

# Jaccard

@inline eval_start(::Jaccard, a::AbstractArray, b::AbstractArray) = 0, 0
@inline function eval_op(::Jaccard, s1, s2)
    denominator = max(s1, s2)
    numerator = min(s1, s2)
    numerator, denominator
end
@inline function eval_reduce(::Jaccard, s1, s2)
    a = s1[1] + s2[1]
    b = s1[2] + s2[2]
    a, b
end
@inline eval_end(::Jaccard, a) = 1 - (a[1]/a[2])
jaccard(a::AbstractArray, b::AbstractArray) = evaluate(Jaccard(), a, b)

# Tanimoto

@inline eval_start(::RogersTanimoto, a::AbstractArray, b::AbstractArray) = 0, 0, 0, 0
@inline function eval_op(::RogersTanimoto, s1, s2)
  tt = s1 && s2
  tf = s1 && !s2
  ft = !s1 && s2
  ff = !s1 && !s2
  tt, tf, ft, ff
end
@inline function eval_reduce(::RogersTanimoto, s1, s2)
    a = s1[1] + s2[1]
    b = s1[2] + s2[2]
    c = s1[3] + s2[3]
    d = s1[4] + s1[4]
    a, b, c, d
end
@inline function eval_end(::RogersTanimoto, a)
    numerator = 2(a[2] + a[3])
    denominator = a[1] + a[4] + 2(a[2] + a[3])
    numerator / denominator
end
rogerstanimoto{T <: Bool}(a::AbstractArray{T}, b::AbstractArray{T}) = evaluate(RogersTanimoto(), a, b)

###########################################################
#
#  Special method
#
###########################################################

# SqEuclidean
function pairwise!(r::AbstractMatrix, dist::SqEuclidean, a::AbstractMatrix, b::AbstractMatrix)
    At_mul_B!(r, a, b)
    sa2 = sumabs2(a, 1)
    sb2 = sumabs2(b, 1)
    pdist!(r, sa2, sb2)
end
function pdist!(r, sa2, sb2)
    for j = 1 : size(r,2)
        sb = sb2[j]
        @simd for i = 1 : size(r,1)
            @inbounds r[i,j] = sa2[i] + sb - 2 * r[i,j]
        end
    end
    r
end
function pairwise!(r::AbstractMatrix, dist::SqEuclidean, a::AbstractMatrix)
    m, n = get_pairwise_dims(r, a)
    At_mul_B!(r, a, a)
    sa2 = sumsq_percol(a)
    for j = 1 : n
        for i = 1 : j-1
            @inbounds r[i,j] = r[j,i]
        end
        @inbounds r[j,j] = 0
        for i = j+1 : n
            @inbounds r[i,j] = sa2[i] + sa2[j] - 2 * r[i,j]
        end
    end
    r
end

# Euclidean
function pairwise!(r::AbstractMatrix, dist::Euclidean, a::AbstractMatrix, b::AbstractMatrix)
    m, na, nb = get_pairwise_dims(r, a, b)
    At_mul_B!(r, a, b)
    sa2 = sumsq_percol(a)
    sb2 = sumsq_percol(b)
    for j = 1 : nb
        for i = 1 : na
            @inbounds v = sa2[i] + sb2[j] - 2 * r[i,j]
            @inbounds r[i,j] = isnan(v) ? NaN : sqrt(max(v, 0.))
        end
    end
    r
end

function pairwise!(r::AbstractMatrix, dist::Euclidean, a::AbstractMatrix)
    m, n = get_pairwise_dims(r, a)
    At_mul_B!(r, a, a)
    sa2 = sumsq_percol(a)
    for j = 1 : n
        for i = 1 : j-1
            @inbounds r[i,j] = r[j,i]
        end
        @inbounds r[j,j] = 0
        for i = j+1 : n
            @inbounds v = sa2[i] + sa2[j] - 2 * r[i,j]
            @inbounds r[i,j] = isnan(v) ? NaN : sqrt(max(v, 0.))
        end
    end
    r
end

# CosineDist

function pairwise!(r::AbstractMatrix, dist::CosineDist, a::AbstractMatrix, b::AbstractMatrix)
    m, na, nb = get_pairwise_dims(r, a, b)
    At_mul_B!(r, a, b)
    ra = sqrt!(sumsq_percol(a))
    rb = sqrt!(sumsq_percol(b))
    for j = 1 : nb
        @simd for i = 1 : na
            @inbounds r[i,j] = max(1 - r[i,j] / (ra[i] * rb[j]), 0)
        end
    end
    r
end
function pairwise!(r::AbstractMatrix, dist::CosineDist, a::AbstractMatrix)
    m, n = get_pairwise_dims(r, a)
    At_mul_B!(r, a, a)
    ra = sqrt!(sumsq_percol(a))
    for j = 1 : n
        @simd for i = j+1 : n
            @inbounds r[i,j] = max(1 - r[i,j] / (ra[i] * ra[j]), 0)
        end
        @inbounds r[j,j] = 0
        for i = 1 : j-1
            @inbounds r[i,j] = r[j,i]
        end
    end
    r
end

# CorrDist
_centralize_colwise(x::AbstractVector) = x .- mean(x)
_centralize_colwise(x::AbstractMatrix) = x .- mean(x, 1)
function colwise!(r::AbstractVector, dist::CorrDist, a::AbstractMatrix, b::AbstractMatrix)
    colwise!(r, CosineDist(), _centralize_colwise(a), _centralize_colwise(b))
end
function colwise!(r::AbstractVector, dist::CorrDist, a::AbstractVector, b::AbstractMatrix)
    colwise!(r, CosineDist(), _centralize_colwise(a), _centralize_colwise(b))
end
function pairwise!(r::AbstractMatrix, dist::CorrDist, a::AbstractMatrix, b::AbstractMatrix)
    pairwise!(r, CosineDist(), _centralize_colwise(a), _centralize_colwise(b))
end
function pairwise!(r::AbstractMatrix, dist::CorrDist, a::AbstractMatrix)
    pairwise!(r, CosineDist(), _centralize_colwise(a))
end


@testset "MPS" begin
    @testset "Inner" begin
        l = 10;
        pchs = U1Charges(-2:2)
        vchs = U1Charges(-3:3)
        pds = [1 for i in pchs]
        vds = [2 for i in vchs]
        mpo = hubbardMPO(1, 1, pchs,l, T = ComplexF64)
        mps = randmps(l, ComplexF64, (U1(), (pchs, vchs), (pds, vds)))
        mpsa = todense(mps)
        mpoa = todense(mpo)

        @test inner(mps) ≈ inner(mpsa)
        @test inner(mps,mpo,mps) ≈ inner(mpsa,mpoa,mpsa)
    end

    @testset "Canonicalization" begin
        l = 10;
        pchs = U1Charges(-2:2)
        vchs = U1Charges(-3:3)
        pds = [1 for i in pchs]
        vds = [2 for i in vchs]
        mps = randmps(l, ComplexF64, (U1(), (pchs, vchs), (pds, vds)))
        for i in 1:l-1
            lmps = MatrixProductStates.leftcanonicalize!(deepcopy(mps),i)
            lenv = foldl(MatrixProductStates.contractsites,zip(lmps[1:i],lmps[1:i]'), init=())
            @test all([v ≈ I for (k,v) in tensor(lenv)])
        end

        for i in l-1:-1:1
            rmps = MatrixProductStates.rightcanonicalize!(deepcopy(mps),i)
            renv = foldr(MatrixProductStates.contractsites,[zip(rmps[i+1:l],rmps[i+1:l]')...], init=())
            @test all([v ≈ I for (k,v) in tensor(renv)])
        end

        for i in 1:l-1
            cmps = canonicalize(mps,i)
            renv = foldr(MatrixProductStates.contractsites,[zip(cmps[i+1:l],cmps[i+1:l]')...], init=())
            x2 = all([v ≈ I for (k,v) in tensor(renv)])
            lenv = foldl(MatrixProductStates.contractsites,[zip(cmps[1:i],cmps[1:i]')...], init=())
            @test all([v ≈ I for (k,v) in tensor(lenv)])
        end
    end
end


### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# ╔═╡ a62a0de6-04b8-11eb-342f-bfea16900b85
begin
	using Pkg
	cd(joinpath(dirname(@__FILE__),".."))
    Pkg.activate(pwd())
	using PlutoUI
	using JuMP, GLPK, Tulip, Optim
	# Non-linear problem
	using Ipopt, Plots
	#using JuMP, Tulip, GLPK, LinearAlgebra
	#using Distributions, Plots, StatsPlots, LaTeXStrings, Measures
end

# ╔═╡ 86702ce8-04b7-11eb-22d1-b7f1437cf740
md"""
# Linear programming II
## Maximize flow in a network

We try to maximize the flow in a network using Linear Programming.

Let $N = (V, E)$ be a directed graph, where $V$ denotes the set of vertices and $E$ is the set of edges. Let $s ∈ V$ and $t ∈ V$ be the source and the sink of $N$, respectively. The capacity of an edge is a mapping $c : E \mapsto \mathbb{R}^+$, denoted by $c_{u,v}$ or $c(u, v)$. It represents the maximum amount of flow that can pass through an edge.

A flow is a mapping $f : E \mapsto \mathbb{R}^+$ , denoted by $f_{uv}$ or  $f(u, v)$, subject to the following two constraints:

* Capacity Constraint:

```math
\forall e \in E: f_{uv} \le c_{uv}
```

* Conservation of Flows:

```math
\forall v \in V\setminus\{s,t\} : \sum_{u:(u,v)\in E}f_{uv} = \sum_{w:(v,w)\in E} f_{vw}
```

We want to maximize the flow in the network, i.e.
```math
\max |f| = \max \sum_{v:(s,v)\in E}f_{sv} = \max \sum_{v:(v,t)\in E}f_{vt}
```

#### Setting:
Consider the following network:

$(PlutoUI.LocalResource("./Exercises/img/network.png"))

We want to:
1. Determine the maximal flow in the network
2. Be able to get a troughput of 35 from the source node to the sink node, whilst keeping the costs limited. Each link has a possible increase, with an associated cost (cf. table)

$(PlutoUI.LocalResource("./Exercises/img/networkcost.png"))
"""

# ╔═╡ caf5c6d8-04b8-11eb-1c04-0966207fbf28
# given set-up
begin
	# Topology and maximum flow matrix
	C = [0 13 6 10 0 0 0;
		 0 0  0  9 5 7 0;
		 0 0  0  8 0 0 0;
		 0 0  0  0 3 0 12;
		 0 0  0  0 0 4 6;
		 0 0  0  0 0 0 9;
		 0 0  0  0 0 0 0;
	];
	# extra capacity
	xcap = [ 0 6  4  3 0 0 0;
			 0 0  0  4 5 3 0;
			 0 0  0  5 0 0 0;
			 0 0  0  0 2 0 5;
			 0 0  0  0 0 2 4;
			 0 0  0  0 0 0 5;
			 0 0  0  0 0 0 0;
	];
	# cost per increased capacity
	xcost= [ 0 2.8  2.5  2.8 0   0   0;
			 0 0    0    2.5 3.1 1.6 0;
			 0 0    0    3.9 0   0   0;
			 0 0    0    0   2.8 0   1.6;
			 0 0    0    0   0   4.6 2.9;
			 0 0    0    0   0   0   1.8;
			 0 0    0    0   0   0   0;
	];
end

# ╔═╡ 5d0d7415-38fe-4c31-9c3d-59b2cbb54a8e
let
	model = Model(Ipopt.Optimizer)
	@variable(model, F[1:7,1:7])
	@constraint(model, F .<= C) # limite capacité
	for i = 2:6
		@constraint(model, sum(F[i,:]) == sum(F[:,i]))
	end
	@constraint(model, F .>= 0)
	@objective(model, Max, sum(F[1,:]))
	optimize!(model)
	abs.(round.(value.(F),digits=6))
end

# ╔═╡ 0f118a6f-9609-44b3-9e08-2e809f516377
let
	model = Model(Ipopt.Optimizer)
	@variable(model,F[1:size(C,1), 1:size(C,2)])
	# targt
	@objective(model, Max, sum(F[1,:])) #@objective(model, Max, sum(F[:,end]))
	# simplifier un peu
	for ind in findall(iszero,C)
		JuMP.fix(F[ind], 0., force=true)
	end
	# constraints
	for i = 2:6
		@constraint(model, sum(F[i,:]) == sum(F[:,i]))
	end
	@constraint(model, F .<= C)
	@constraint(model, F .>= 0)
	#
	optimize!(model)
	value.(F)
	# actuellement débit = 27
end

# ╔═╡ e6ba22db-4fc3-4a23-bb56-fa500ef616d2
let
	model = Model(Ipopt.Optimizer)
	@variable(model,F[1:size(C,1), 1:size(C,2)])
	# targt
	@objective(model, Min, sum(F .* xcost)) # faux raisonnement :-(
	# simplifier un peu
	for ind in findall(iszero,C)
		JuMP.fix(F[ind], 0., force=true)
	end
	# constraints
	for i = 2:6
		@constraint(model, sum(F[i,:]) == sum(F[:,i]))
	end
	@constraint(model, F .<= C .+ xcap)
	@constraint(model, F .>= 0)
	@constraint(model, sum(F[1,:]) == 35.)
	#
	optimize!(model)
	value.(F)
	# actuellement débit = 27
	Δ_cap = value.(F) - C
	ind_pos = Δ_cap .> 0
	sum((Δ_cap .* xcost)[ind_pos])
	Δ_cap
end

# ╔═╡ c38feffb-245b-4812-bcec-bdecd43d6c6c
let
	model = Model(Ipopt.Optimizer)
	@variable(model,F[1:size(C,1), 1:size(C,2)]) # capacité existante
	@variable(model,G[1:size(C,1), 1:size(C,2)]) # capacité supplémentaire
	# target
	@objective(model, Min, sum(G .* xcost)) # bon raisonnement :-)
	# simplifier un peu
	for ind in findall(iszero,C)
		JuMP.fix(F[ind], 0., force=true)
		JuMP.fix(G[ind], 0., force=true)
	end
	# constraints
	for i = 2:6
		@constraint(model, sum(F[i,:] + G[i,:]) == sum(F[:,i] + G[:,i]))
	end
	@constraint(model, F .<= C)
	@constraint(model, F .>= 0)
	@constraint(model, G .<= xcap)
	@constraint(model, G .>= 0)
	@constraint(model, sum(F[1,:]+ G[1,:]) == 35.)
	#
	optimize!(model)
	value.(F)
	# actuellement débit = 27
	sum(value.(G) .* xcost)
	value.(G)
end

# ╔═╡ d82dbb55-1d31-445a-b5e4-ec11ec4fcf65

# ╔═╡ 84804f16-8dfb-4bad-afe7-fe8ccd854e39

# ╔═╡ d4f240a2-73c5-4828-b728-d922db6148b8

# ╔═╡ 27080e2e-a7df-4acb-93a9-98f2352fd8b5

# ╔═╡ 508ae591-46f9-4dd9-97b8-08a5fd90e6e1

# ╔═╡ d3dcd770-058d-11eb-2829-fbb18e85f401
begin
	model = Model(Ipopt.Optimizer)
	@variable(model, F[1:size(C,1),1:size(C,2)])
	# gekende waarde op nul zetten
	for ind in findall(iszero, C)
		JuMP.fix(F[ind],0.)
	end
	@objective(model, Max, sum(F[1,:])) # of ook @objective(model, Max, sum(F[:,end]))
	# capacity constraint
	@constraint(model, F .<= C)
	# maintain flow
	for i = 2:6
		@constraint(model, sum(F[:,i]) == sum(F[i,:]))
	end

	optimize!(model)
end

# ╔═╡ c78e64e1-df73-4c6d-b705-2e1766d8ef3a
let
	# capaciteitsverhoging
	model = Model(Ipopt.Optimizer)
	@variable(model, F[1:size(C,1),1:size(C,2)])
	# gekende waarde op nul zetten
	for ind in findall(iszero, C)
		JuMP.fix(F[ind],0.)
	end
	# minimise total cost
	@objective(model, Min, sum(F .* xcost) ) # hier zit een redeneringsprobleem (!)
	# better?
	#helpfun(F) = F .- C .> 0
	#@objective(model, Min, filter(x-> x>0, (F .- C) .* xcost))

	# capacity constraint
	@constraint(model, F .<= C .+ xcap)
	@constraint(model, F .>= 0)
	# maintain flow
	for i = 2:6
		@constraint(model, sum(F[:,i]) == sum(F[i,:]))
	end
	# demand specifc throughput
	@constraint(model, sum(F[1,:]) == 35)

	optimize!(model)
	value.(F)
	# value of extra cost?
	Δ_cap = value.(F) .- C
	indpos = Δ_cap .> 0.
	sum(Δ_cap[indpos] .* xcost[indpos])
	#helpfun(F)
end

# ╔═╡ 74a20456-a91a-44de-9e4b-3e6d86440392
let
	# capaciteitsverhoging
	model = Model(Ipopt.Optimizer)
	@variable(model, F[1:size(C,1),1:size(C,2)]) # oorspronkelijk
	@variable(model, G[1:size(C,1),1:size(C,2)]) # extra
	# gekende waarde op nul zetten
	for ind in findall(iszero, C)
		JuMP.fix(F[ind],0.)
		JuMP.fix(G[ind],0.)
	end
	# minimise total cost of additional capacity
	@objective(model, Min, sum(G .* xcost) )

	# capacity constraint
	@constraint(model, F .<= C)
	@constraint(model, G .<= xcap)
	@constraint(model, F .>= 0)
	@constraint(model, G .>= 0)
	# maintain flow
	for i = 2:6
		@constraint(model, sum(F[:,i] .+ G[:,i]) == sum(F[i,:] .+ G[i,:]))
	end
	# demand specifc throughput
	@constraint(model, sum(F[1,:] .+ G[1,:]) == 35)

	optimize!(model)
	value.(F) .+ value.(G)
	# total cost:
	sum(value.(G) .* xcost)
end

# ╔═╡ 6ea03142-b487-4459-8391-cd50dae56340

# ╔═╡ 42371ec0-55fa-459b-bcb2-a6ea4363d876
value.(F)

# ╔═╡ 0a95f12e-c5eb-490a-bc7c-0cd7a70edddb
sum(value.(F)[1,:])

# ╔═╡ ef581235-7b51-4d2d-8097-7f6637c40b24
let
	model = Model(Ipopt.Optimizer)
	@variable(model, F[1:size(C,1),1:size(C,2)] >=0)
	# gekende waarde op nul zetten
	for ind in findall(iszero, C)
		JuMP.fix(F[ind],0, force=true)
	end
	@objective(model, Max, sum(F[1,:]))
	# capacity constraint
	@constraint(model, F .<= C .+ xcap)
	# maintain flow
	for i = 2:6
		@constraint(model, sum(F[:,i]) == sum(F[i,:]))
	end

	optimize!(model)
	value.(F)
end

# ╔═╡ 1283cdae-1a6b-4a4e-be2b-a9072bfc3972

# ╔═╡ 55ac8bb2-eb86-4536-9a89-245cf3883763

# ╔═╡ 7e852fdd-4363-4e15-ab24-e76448fbe1cf

# ╔═╡ a510792e-04c0-11eb-0f6e-a7e40dfe602c
md"""
## Optimizing an investment portfolio

In 1952 [Harry Max Markowitz](https://en.wikipedia.org/wiki/Harry_Markowitz) proposed a new approach for the optimization of an investment portfolio. This ultimately led to him winning the Nobel Prize in Economic Sciences in 1990. The idea is relatively simple:

Given a portfolio with $n$ stock proportions $S_1,S_2,\dots, S_n$, we want to maximize the return (=profit) and minimize the risk. The goal is to find the values $S_i$ that lead to either a minimum risk attribution with a minimal return or that lead to a maximum return attribution with a maximal risk.

Remembering that $\sigma^{2}_{\sum_{i=1}^{n}X_i}= \sum_{i=1}^{n}\sigma^2_{X_i} + \sum_{i \ne j}\text{Cov}(X_i,X_j) $, the risk can be expressed in terms of the covariance matrix $\Sigma$:

$$S^\mathsf{T} \Sigma S $$

The return can be expressed as:
$$\mu^\mathsf{T}S$$

Consider the following portfolio problem:
You are given the covariance matrix and expected returns and you want study several approaches. For each case you should formulate a proper Linear/Quadratic Programming form.
1. Ignore the risk and go for optimal investment (i.e. maximal return)
2. Same as (1), but a single stock can be at most 40% of the portfolio
3. Minimize the risk, with a lower bound on the return e.g. with at least 35% expected return
4. Make a graph for:
    * the minimal risk in fuction of the expected return.
    * the distribution of the portfolio with the minimal risk in function of the expected return
    * the final portfolio value in function of the expected return
"""

# ╔═╡ c4eca510-04c0-11eb-0e29-a5dcc1386bd5
# data for problem
begin
	P = [60; 127; 4; 50; 150; 20] # stock prices
	μ = [0.2; 0.42; 1.; 0.5; 0.46; 0.3] # expected returns
	Σ = [0.032 0.005 0.03 -0.031 -0.027 0.01;
		 0.005 0.1 0.085 -0.07 -0.05 0.02;
		 0.03 0.085 1/3 -0.11 -0.02 0.042;
		 -0.031 -0.07 -0.11 0.125 0.05 -0.06;
		 -0.027 -0.05 -0.02 0.05 0.065 -0.02;
		 0.01 0.02 0.042 -0.06 -0.02 0.08]; # covariance matrix
end

# ╔═╡ 0589b400-06a9-4f91-99aa-97e52f191dd0
let
	# Q1
	model = Model(Ipopt.Optimizer)
	# variables
	@variable(model,S[1:length(μ)]) # % van ieder aandeel in mijn portefeuille
	# objective
	@objective(model, Max, μ'*S) # max profit, risico bestaat niet
	# constraints
	@constraint(model, sum(S) == 1.) # total 100 %
	@constraint(model, S .>= 0.) # geen negatieve bijdrages
	# los het op
	optimize!(model)
	value.(S)

	# Q2
	model = Model(Ipopt.Optimizer)
	# variables
	@variable(model,S[1:length(μ)]) # % van ieder aandeel in mijn portefeuille
	# objective
	@objective(model, Max, μ'*S) # max profit, risico bestaat niet
	# constraints
	@constraint(model, sum(S) == 1.) # total 100 %
	@constraint(model, S .>= 0.) # geen negatieve bijdrages
	@constraint(model, S .<= 0.4) # diversificatie!
	# los het op
	optimize!(model)
	value.(S)

	# Q3
	model = Model(Ipopt.Optimizer)
	# variables
	@variable(model,S[1:length(μ)]) # % van ieder aandeel in mijn portefeuille
	# objective
	@objective(model, Min, S'*Σ*S) # min risk
	# constraints
	@constraint(model, μ'*S >= 0.35) # min return van 0.35
	@constraint(model, sum(S) == 1.) # total 100 %
	@constraint(model, S .>= 0.) # geen negatieve bijdrages
	@constraint(model, S .<= 0.4) # diversificatie!
	# los het op
	optimize!(model)
	value.(S), μ'*value.(S)
	termination_status(model)
end

# ╔═╡ e15f2647-04e0-4882-a193-5a63813852d4
let
	# ignorer le risque & maximiser le gain
	model = Model(Ipopt.Optimizer)
	@variable(model, S[1:length(P)])
	@constraint(model, S .<= 1.) # pas plus de 100% par action
	@constraint(model, S .>= 0.) # pas de proportion négative
	@constraint(model, sum(S) == 1.) # Σproportion = 1
	@objective(model, Max, μ' * S)
	optimize!(model)
	value.(S)
	# ignorer le risque & maximiser le gain & <=40% par action
	model = Model(Ipopt.Optimizer)
	@variable(model, S[1:length(P)])
	@constraint(model, S .<= 0.4) # pas plus de 100% par action
	@constraint(model, S .>= 0.) # pas de proportion négative
	@constraint(model, sum(S) == 1.) # Σproportion = 1
	@objective(model, Max, μ' * S)
	optimize!(model)
	value.(S)
	# minimiser le risque & gain minimal de 0.35
	model = Model(Ipopt.Optimizer)
	@variable(model, S[1:length(P)])
	@constraint(model, S .<= 0.4) # pas plus de 100% par action
	@constraint(model, S .>= 0.) # pas de proportion négative
	@constraint(model, sum(S) == 1.) # Σproportion = 1
	@constraint(model, μ' * S >= 0.35) # gain minimal à assurer
	@objective(model, Min, S' * Σ * S)

	optimize!(model)
	value.(S)

end

# ╔═╡ c6cb56d8-058d-11eb-2f71-73a5e5ca3dd0

# ╔═╡ ac5e426e-55f6-4ed0-98e4-19fb00ff1aec
let
	model = Model(Tulip.Optimizer)
	N = 20
	@variable(model, T[1:N] >= 0) # theorie part
	@variable(model, A[1:N] >= 0) # application part
	@constraint(model, T .<= 1) # logic
	@constraint(model, A .<= 1) # logic
	# teaching model
	a = 2; b = 3;
	@constraint(model, sum(A[1:N]) <= sum( T[b+1:N]))
	for i = 1:N
		if i <= b
			JuMP.fix(T[i], 1., force=true)
			JuMP.fix(A[i], 0., force=true)
		else
			@constraint(model, T[i] + A[i] == 1.)
		end
	end
	# emotional states
	s₀ = 0.
	α = -0.1; θ = 0.05; β = 1.4;
	S = [(1-θ)*s₀ + θ*(α*T[1] + β*A[1])]
	for i = 2:N
		push!(S, (1-θ)*S[i-1] + θ*(α*T[i-1] + β*A[i-1]))
	end
	@objective(model, Max, S[end])
	optimize!(model)
	termination_status(model)
	plot(value.(T))
	plot!(value.(A))
end

# ╔═╡ cb6898d6-0565-11eb-2a3c-1514a0fa4f50
md"""
## Optimal course planning
Suppose a professor teaches a course with $N=20$ lectures. We must decide how to split each lecture between theory and applications. Let $T_i$ and $A_i$ denote the fraction of the i$^{\text{th}}$ lecture devoted to theory and applications, for $i=1,\dots,N$. We can already determine the following:

```math
\forall i: T_i \ge 0, A_i \ge 0, T_i+A_i =1.
```

As you may know from experience, you need to cover a certain amount of theory before you can start doing applications. For this application consider the following model:

$$\sum_{i=1}^{N} A_i \le \phi \left( \sum_{i=1}^{N} T_i \right)$$

We interpret $\phi(u)$ as the cumulative amount of applications that can be covered, when the cumulative amount of theory covered is $u$. We will use the simple form $\phi(u) = a(u − b)$, with $a=2, b=3$, which means that no applications can be covered until $b$ lectures of the theory are covered; after that, each lecture of theory covered opens the possibility of covering a lecture on applications.

Psychological studies have shown that the theory-applications split affects the emotional state of students differently. Let $s_i$ denote the emotional state of a student after lecture $i$, with $s_i = 0$ meaning neutral, $s_i > 0$ meaning happy, and $s_i < 0$ meaning unhappy. Careful studies have shown that $s_i$ evolves via a linear recursion dynamic:

$$s_i =(1−\theta)s_{i−1} +\theta(\alpha T_i +\beta A_i)\text{ with }\theta \in[0,1]$$

with $s_0=0$. In order to make sure that the student leave with a good feeling at the end of the course, we try to maximize $s_N$, i.e. the emotional state after the last lecture.

Questions:
1. Determine the optimal split that leads to the most positive emotional state (for $\theta = 0.05, \alpha = -0.1, \beta = 1.4$);
2. Show the course repartition graphically
3. Determine values for $\alpha$ and $\beta$ that lead to a neutral result at the end of the course. Can you give an interpretation to these values?
"""

# ╔═╡ b76bae36-058d-11eb-216c-dfb8abd85772

# ╔═╡ Cell order:
# ╠═a62a0de6-04b8-11eb-342f-bfea16900b85
# ╠═5d0d7415-38fe-4c31-9c3d-59b2cbb54a8e
# ╟─86702ce8-04b7-11eb-22d1-b7f1437cf740
# ╠═0f118a6f-9609-44b3-9e08-2e809f516377
# ╠═e6ba22db-4fc3-4a23-bb56-fa500ef616d2
# ╠═c38feffb-245b-4812-bcec-bdecd43d6c6c
# ╠═caf5c6d8-04b8-11eb-1c04-0966207fbf28
# ╠═d82dbb55-1d31-445a-b5e4-ec11ec4fcf65
# ╠═84804f16-8dfb-4bad-afe7-fe8ccd854e39
# ╠═d4f240a2-73c5-4828-b728-d922db6148b8
# ╠═27080e2e-a7df-4acb-93a9-98f2352fd8b5
# ╠═508ae591-46f9-4dd9-97b8-08a5fd90e6e1
# ╠═d3dcd770-058d-11eb-2829-fbb18e85f401
# ╠═c78e64e1-df73-4c6d-b705-2e1766d8ef3a
# ╠═74a20456-a91a-44de-9e4b-3e6d86440392
# ╠═6ea03142-b487-4459-8391-cd50dae56340
# ╠═42371ec0-55fa-459b-bcb2-a6ea4363d876
# ╠═0a95f12e-c5eb-490a-bc7c-0cd7a70edddb
# ╠═ef581235-7b51-4d2d-8097-7f6637c40b24
# ╠═1283cdae-1a6b-4a4e-be2b-a9072bfc3972
# ╠═55ac8bb2-eb86-4536-9a89-245cf3883763
# ╠═7e852fdd-4363-4e15-ab24-e76448fbe1cf
# ╟─a510792e-04c0-11eb-0f6e-a7e40dfe602c
# ╠═0589b400-06a9-4f91-99aa-97e52f191dd0
# ╠═c4eca510-04c0-11eb-0e29-a5dcc1386bd5
# ╠═e15f2647-04e0-4882-a193-5a63813852d4
# ╠═c6cb56d8-058d-11eb-2f71-73a5e5ca3dd0
# ╠═ac5e426e-55f6-4ed0-98e4-19fb00ff1aec
# ╟─cb6898d6-0565-11eb-2a3c-1514a0fa4f50
# ╠═b76bae36-058d-11eb-216c-dfb8abd85772

Exercises/PS05 - Linear programming II.jl


include("utils_dg.jl")

K = 6  # Number of elements

n = 4
nn = 1
r = jacobiGL(0, 0, n)
V = vandermonde(r, 0, 0, n)
fr = collect(range(-1,1,length=100))
fV = vandermonde(fr, 0, 0, n)
refine = fV * inv(V)
# cell average
avg = Diagonal(zeros(n+1))
avg[1:nn] .= 1
p0 = V * avg * inv(V)
# polyone
nn = 2
tmp = zeros(n+1)
tmp[1:nn] .= 1
avg = Diagonal(tmp)
p1 = V * avg * inv(V)
# poly two
nn = 2+1
tmp = zeros(n+1)
tmp[1:nn] .= 1
avg = Diagonal(tmp)
p2 = V * avg * inv(V)

# poly 6
Ω = (; a=0, b=2π, periodic = true)
a = Ω.a
b = Ω.b
mesh = Mesh(K, n, Ω.a, Ω.b, periodic = Ω.periodic)
x = mesh.x
weights = sum(mesh.M, dims = 1)
Δx = x[2] - x[1]

# inexact integration
DM = Diagonal(sum(mesh.M, dims = 1)[:])
mesh.M .= DM
mesh.Mi .= inv(DM)
mesh.lift[:,1] .= mesh.Mi[1,:]
mesh.lift[:,end] .= mesh.Mi[end,:]

b = Ω.b
a = Ω.a

u_first  = mesh.D *  @.sin( 2 * 2π/(b-a) * x)
using Random
Random.seed!(123)
u_first = compute_volume_terms(u_first, mesh) + 0.1 * compute_surface_terms(randn(size(u_first)), mesh)
uo = refine * u_first
x = refine * x

DVu = refine * compute_volume_terms(u_first, mesh)
DAu = refine * compute_surface_terms(u_first, mesh)
Du = DVu + DAu

fig = Figure(resolution = (1700+600, 1000+400))

lims = (extrema(x)..., extrema(uo)...)
axo = Axis(fig[1,1], xlabel = L"x", ylabel = L"y", ylabelsize = 22,
    xlabelsize= 22, xgridstyle=:dash, ygridstyle=:dash, xtickalign = 1,
    xticksize=10, ytickalign=1, yticksize=10,  xlabelpadding = -10)

axo.limits = (extrema(x)..., (extrema(uo) .* 1.1)...)
axo.title = "Original"
axo.titlesize = 32
for i in 1:K
    lines!(axo, x[:,i], uo[:,i], linewidth = 5)
end

ax0 = Axis(fig[1,2], xlabel = L"x", ylabel = L"y", ylabelsize = 22,
    xlabelsize= 22, xgridstyle=:dash, ygridstyle=:dash, xtickalign = 1,
    xticksize=10, ytickalign=1, yticksize=10,  xlabelpadding = -10, ylims = (-1.1, 1.1))
lims2 = (extrema(x)..., (extrema(DVu) .* 1.1)...)
ax0.limits = lims2
ax0.title = "Volume Contribution"
ax0.titlesize = 32
for i in 1:K
    lines!(ax0, x[:,i], DVu[:,i], linewidth = 5)
end

ax1 = Axis(fig[2,1], xlabel = L"x", ylabel = L"y", ylabelsize = 22,
    xlabelsize= 22, xgridstyle=:dash, ygridstyle=:dash, xtickalign = 1,
    xticksize=10, ytickalign=1, yticksize=10,  xlabelpadding = -10, ylims = (-1.1, 1.1))
ax1.limits = (extrema(x)..., (extrema(DAu) .* 1.1)...)
ax1.title = "Surface Contribution"
ax1.titlesize = 32
for i in 1:K
    lines!(ax1, x[:,i], DAu[:,i], linewidth = 5)
end

ax3 = Axis(fig[2,2], xlabel = L"x", ylabel = L"y", ylabelsize = 22,
    xlabelsize= 22, xgridstyle=:dash, ygridstyle=:dash, xtickalign = 1,
    xticksize=10, ytickalign=1, yticksize=10,  xlabelpadding = -10, ylims = (-1.1, 1.1))
ax3.limits = (extrema(x)..., (extrema(Du) .* 1.1)...)
ax3.title = "Derivative"
ax3.titlesize = 32
for i in 1:K
    lines!(ax3, x[:,i], Du[:,i], linewidth = 5)
end


# %%
if !("." in LOAD_PATH) push!(LOAD_PATH, ".") end
using Revise
@time using CellularBase
using BenchmarkTools
using Random
# %%
mutable struct SandPile <: AbstractGrid{Int,2}
    state::Array{Int,2}
    otherstate::Array{Int,2}
    h_c::Int
    neighborhood::AbstractNeighborhood
    bc::BoundaryCondition
    function SandPile(state::Array{Int,2}, h_c::Int)
        new(state, similar(state), h_c, VonNeumannNeighborhood{2}(1), FixedMin())
    end
end
@inline iscritical(c::Int, g::SandPile) = c >= g.h_c
@inline iscritical(g::SandPile) = i -> iscritical(i, g)

CellularBase.newstate(s::SandPile) = s.otherstate
function CellularBase.state!(s::SandPile, newstate)
    temp = s.state
    s.state = s.otherstate
    s.otherstate = temp
end

function (_g::SandPile)(current::Int, neighbors::Array{Int, 1})
    h = iscritical(_g)
    return current + h(current)*-4 + count(h, neighbors)
end

CellularBase.possible_states(g::SandPile) = 0:(g.h_c + 3)
# %%

function getindexbc(grid, i)
    if checkbounds(Bool, grid, i)
        @inbounds return grid[i]
    else
        return 0
    end
end
getindexbc(grid) = i -> getindexbc(grid)
function evolveat(grid, hood, index, hc)
    res = grid[index]
    if res >= hc
        res -= 4
    end
    c = count(h -> hc <= getindexbc(grid, h + index), hood)
    res = res + c
    return res
end
function evolve(grid::Array{Int}, hood::Array{CartesianIndex{2},1}, hc::Int; kwargs...)
    simgrid = similar(grid)
    for index in CartesianIndices(grid)
        simgrid[index] = evolveat(grid, hood, index, hc)
    end
    return simgrid
end

function baresimulate(grid::Array{Int}, hc::Int, steps::Int)
    hood = CartesianIndex.([(-1, 0), (1, 0), (0, -1), (0, 1)])
    for i in 1:steps
        grid = evolve(grid, hood, hc)
    end
    return grid
end
function baresimulate!(cbgrid::SandPile, steps)
    for i in 1:steps
        evolve!(cbgrid)
    end
end

function gridtest(cbgrid::SandPile)
    for i = 1:400
        cbgrid(1, [1, 1, 1, 1])
    end
end
function gridtest(ngrid::Array{Int})
    hood = cartesians(VonNeumannNeighborhood{2}(1));
    for i in 1:400
        evolveat(ngrid, hood, CartesianIndex(1,1), 4)
    end
end

CellularBase.boundary_conditions(p::SandPile)::FixedMin = FixedMin()

# %%
println("===============")
Random.seed!(0);
hc = 4;
ngrid = rand(1:hc, 20, 20);
cbgrid = SandPile(copy(ngrid), 4);
hood = cartesians(VonNeumannNeighborhood{2}(1));

@btime baresimulate(ngrid, 4, 5)
# @btime evolve(ngrid, hood, 4)
# @btime evolve(cbgrid)
# @btime CellularBase.evolve!(cbgrid)
@btime simulate!(cbgrid, 5, store_results=false)
# @btime baresimulate!(cbgrid, 1)

# %%


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

using BAT
using Test

using Random, Distributions, StatsBase

@testset "bat_sample" begin
    @testset "IIDSampling" begin
        dist = MvNormal([0.4, 0.6], [2.0 1.2; 1.2 3.0])

        @test length(@inferred(bat_sample(dist, IIDSampling(nsamples = 10^3))).result) == 10^3

        @test @inferred(bat_sample(Random.GLOBAL_RNG, dist)).result isa DensitySampleVector
        @test @inferred(bat_sample(dist, BAT.IIDSampling())).result isa DensitySampleVector
        @test @inferred(bat_sample(Random.GLOBAL_RNG, dist, BAT.IIDSampling())).result isa DensitySampleVector

        samples = @inferred(bat_sample(dist, IIDSampling(nsamples = 10^5))).result
        @test isapprox(mean(samples.v), [0.4, 0.6]; rtol = 0.05)
        @test isapprox(cov(samples.v), [2.0 1.2; 1.2 3.0]; rtol = 0.05)
        @test all(isequal(1), samples.weight)

        resamples = @inferred(bat_sample(samples, OrderedResampling(nsamples = length(samples)))).result
        @test samples == resamples

        dist_bmode = @inferred(bat_findmode(dist)).result
        @test @inferred(length(dist_bmode)) == 2

        dist_sample_vector_bmode = @inferred(bat_findmode(samples)).result
        @test @inferred(length(dist_sample_vector_bmode)) == 2

        isapprox(var(bat_sample(Normal(), BAT.IIDSampling(nsamples = 10^3)).result), [1], rtol = 10^-1)
    end
end


###################################
### Last Update: 29 Oct 20
###################################
### ''Simulation of Optimal Proposal''
### Code follows method/notation in Dietrich & Newsam (96)
###################################

using Distributions
using LinearAlgebra
using FFTW
using Random

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

Random.seed!(123)

#### Parameters for Covariance Function
struct f_th
    s::Float64
    l::Float64
end

#### Parameters for Grid
struct f_gr
    p::Int ## Number of Grid Points Along x-axis is (p+1)
    q::Int ## Number of Grid Points Along y-axis is (q+1)
    P::Float64
    Q::Float64 ## 2D-Field is rectangle [0,P]x[0,Q]
end

#### Parameters for Stations
struct f_st
  st_ij::Array{Int64,2} ## the position of the stations on the grid
                        ## i.e station positions determined using 2 co-ordinates
   st_p::Array{Int64,1} ## the no of grid point occupied by each station
                        ## i.e station positions determined using 1 co-ordinate
end

####################################################################
#### Example #######################################################

th = f_th(1.0, 1.0)

gr = f_gr(20, 20, 40.0, 40.0)

###################################################
#### Example ######################################

p = gr.p; q = gr.q; P = gr.P; Q = gr.Q

Dx = P/p; Dy = Q/q ## Grid Steps Along x and y Directions

###################################################
#### Example ######################################

#### In real case Stations will be read from a file containing co-ordinates;
#### here, I delect position of n2 Stations at random from all grid nodes

n2 = 3; n1 = (p+1)*(q+1)

grid_F = fill(0,(n1,2))

#### the loop contructs 2-dim co-ordinates for all grid points
for j in 0:q

  for i in 0:p

    grid_F[j*(p+1)+i+1,:] = [i, j]

  end

end
####

st_p = sample(1:n1,n2,replace=false)

st_ij = grid_F[st_p,:]

st = f_st(st_ij, st_p)

###################################################
#### Example ######################################

#### I simulate some arbitrary values to be used as observations of height
#### I give an example value to the standard deviation Sobs;
#### I only consider a single time instance

Yobs_t = rand(Normal(0,1), n2)
#### the ordering in the observations is in tandem with the ordering
#### of the stations in variable st

Sobs = 0.01

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

#### Example Specification of Covariance Function
function R(x::Float64, y::Float64, th)

  Cov = (th.s^2)*exp(-(abs(x)+abs(y))/(2.0*th.l))

  return(Cov)

end

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

#### Specification of 'Extented' Covariance Function
#### following Method/Notation in Dietrich & Newsam (96)
function bar_R(x::Float64, y::Float64, th, gr)

  P = gr.P; Q = gr.Q

  bar_Cov = 0.0

  if ( (0<=x) & (x<=P) & (0<=y) & (y<=Q) )

    bar_Cov = R(x, y, th)

  elseif ( (P<=x) & (x<=2*P) & (0<=y) & (y<=Q) )

    bar_Cov = R(2*P-x, y, th)

  elseif ( (0<=x) & (x<=P) & (Q<=y) & (y<=2*Q) )

    bar_Cov = R(x, 2*Q-y, th)

  elseif ( (P<=x) & (x<=2*P) & (Q<=y) & (y<=2*Q) )

    bar_Cov = R(2*P-x, 2*Q-y, th)

  end

  return(bar_Cov)

end

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

#### Obtain (4pq)x(1) Vector bar_rho
#### following Method/Notation in Dietrich & Newsam (96)
function f_bar_rho(th, gr)

  p = gr.p; q = gr.q; P = gr.P; Q = gr.Q

  Dx = P/p; Dy = Q/q

  bar_rho = zeros(4*p*q)

  #########

  for j in 0:(2*q-1)

    for i in 0:(2*p-1)

      bar_rho[(j*2*p)+(i+1)] = bar_R(i*Dx, j*Dy, th, gr)

    end

  end

  #########

  return(bar_rho)

end

###################################################
#### Example ######################################

####bar_rho = f_bar_rho(th, gr)

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

#### Construct W Operation on (4pq)x1 Vector V
#### W is 2D Discrete Fourier Transform
#### following Method/Notation in Dietrich & Newsam (96);
#### V is (4pq)x1; result is (4pq)x1
function f_W(V, gr)

  p = gr.p; q = gr.q

  WV = zeros(4*p*q)

  #### Re-arrange (4pq)x(1) Vector V into an (2p)x(2q) matrix
  #### by using first 2p-elements of V as 1st column,
  #### next 2p-elements as 2nd column, and so on.
  bar_E = reshape(V, (2*p,2*q))

  #### Use Julia's 2D-FFT tranform
  E = sqrt(4*p*q)*ifft(bar_E)
  #### E is in general a Complex-Value (2p)x(2q) Matrix

  #### Re-arrange (2p)x(2q) matrix E
  #### into an (4pq)x1 matrix, by using 1st column of E,
  #### then 2nd column of E, and so on.
  WV = reshape(E,(4*p*q))

  return(WV)

end

###################################################
#### Example ######################################

####[bar_rho f_W(bar_rho, gr)]

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

#### Construct WH Operation on (4pq)x1 Vector V
#### WH is 2D Inverse Discrete Fourier Transform
#### following Method/Notation in Dietrich & Newsam (96);
#### V is (4pq)x1; result is (4pq)x1
function f_WH(V, gr)

  p = gr.p; q = gr.q

  WHV = zeros(4*p*q)

  #### Re-arrange (4pq)x(1) Vector V into an (2p)x(2q) matrix
  #### by using first 2p-elements of V as 1st column,
  #### next 2p-elements as 2nd column, and so on.
  E = reshape(V, (2*p,2*q))

  #### Use Julia's 2D-IFFT tranform
  bar_E = (1.0/sqrt(4*p*q))*fft(E)
  #### bar_E is in general a Complex-Value (2p)x(2q) Matrix

  #### Re-arrange (2p)x(2q) matrix bar_E
  #### into an (4pq)x1 matrix, by using 1st column of bar_E,
  #### then 2nd column of bar_E, and so on.
  WHV = reshape(bar_E,(4*p*q))

  return(WHV)

end

###################################################
#### Example ######################################

####V = bar_rho; WV = f_W(V,gr); WHWV = f_WH(WV,gr);

#### CHECK!! ######################################

####[V WV WHWV] ## Since WH*W = I, for any V, it should be V = WHWV!!

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

###################################################
###################################################
####
#### Construct Covariance Matrix on Original Grid
#### following Method/Notation in Dietrich & Newsam (96)
#### This Function is Needed only for Testing!!
####function f_R11(th, gr)
####
####  p = gr.p; q = gr.q; P = gr.P; Q = gr.Q
####
####  Dx = P/p; Dy = Q/q; n1 = (p+1)*(q+1)
####
####  R11 = zeros(n1, n1)
####
####  for m in 0:q
####
####    for l in 0:p
####
####      for j in 0:q
####
####        for i in 0:p
####
####          d1 = m*(p+1)+l+1; d2 = j*(p+1)+i+1
####
####          R11[d1,d2] = R(abs(l*Dx-i*Dx), abs(m*Dy-j*Dy), th)
####
####        end
####
####      end
####
####    end
####
####  end
####
####  return(R11)
####
####end
####
###################################################
#### Example ######################################
####
####R11 = f_R11(th, gr)
####
#### Construct Covariance Matrix on Extended Grid
#### following Method/Notation in Dietrich & Newsam (96)
#### This Function is Needed only for Testing!!
####function f_bar_R11(th, gr)
####
####  p = gr.p; q = gr.q; P = gr.P; Q = gr.Q
####
####  Dx = P/p; Dy = Q/q; bar_n1 = 4*p*q
####
####  bar_R11 = zeros(bar_n1, bar_n1)
####
####  for m in 0:(2*q-1)
####
####    for l in 0:(2*p-1)
####
####      for j in 0:(2*q-1)
####
####        for i in 0:(2*p-1)
####
####          d1 = m*2*p+l+1; d2 = j*2*p+i+1
####
####          bar_R11[d1,d2] = bar_R(abs(l*Dx-i*Dx), abs(m*Dy-j*Dy), th, gr)
####
####        end
####
####      end
####
####    end
####
####  end
####
####  return(bar_R11)
####
####end
####
###################################################
#### Example ######################################
####
####bar_R11 = f_bar_R11(th, gr)
####
###################################################
###################################################

#### Construct Covariance between Stations and Extended Grid
#### following Method/Notation in Dietrich & Newsam (96);
#### result is matrix of dimnesions n2xbar_n1
function f_bar_R21(th, st, gr)

  st_ij = st.st_ij; n2 = size(st_ij)[1];
  ## No of Stations

  p = gr.p; q = gr.q; P = gr.P; Q = gr.Q

  Dx = P/p; Dy = Q/q

  bar_R21 = zeros(n2, 4*p*q)

  for m in 1:n2

    for j in 0:(2*q-1)

      for i in 0:(2*p-1)

        d2 = j*2*p+i+1

        bar_R21[m,d2] = bar_R(abs(st_ij[m,1]-i)*Dx, abs(st_ij[m,2]-j)*Dy, th, gr)

      end

    end

  end

  return(bar_R21)

end

###################################################
#### Example ######################################

####bar_R21 = f_bar_R21(th, st, gr)

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

#### Construct Covariance between Observations
#### following Method/Notation in Dietrich & Newsam (96);
#### result is matrix of dimnesions n2xn2
function f_R22(th, st, Sobs, gr)

  st_ij = st.st_ij; n2 = size(st_ij)[1]
  ## No of Stations

  p = gr.p; q = gr.q; P = gr.P; Q = gr.Q

  Dx = P/p; Dy = Q/q

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

  R22 = zeros(n2, n2)

  for i in 1:n2

    xi = st_ij[i,1]
    yi = st_ij[i,2]

    for j in 1:n2

      xj = st_ij[j,1]
      yj = st_ij[j,2]

      R22[i,j] = R(abs(xj-xi)*Dx,abs(yj-yi)*Dy,th)

      if (i==j)

        R22[i,j] = R22[i,j] + Sobs^2

      end

    end

  end

  return(R22)

end

###################################################
#### Example ######################################

####R22 = f_R22(th, st, Sobs, gr)

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

#### Construct Covariance between Original Grid & Station Positions
#### following Method/Notation in Dietrich & Newsam (96)
#### result is matrix of dimnesions n1xn2
function f_R12(th, st, gr)

  st_ij = st.st_ij; n2 = size(st_ij)[1]
  ## No of Stations

  p = gr.p; q = gr.q; P = gr.P; Q = gr.Q

  Dx = P/p; Dy = Q/q; n1 = (p+1)*(q+1)

  R12 = zeros(n1, n2)

  for m in 1:n2

    for j in 0:q

      for i in 0:p

        d1 = j*(p+1)+i+1

        R12[d1,m] = R(abs(st_ij[m,1]-i)*Dx, abs(st_ij[m,2]-j)*Dy, th)

      end

    end

  end

  return(R12)

end

###################################################
#### Example ######################################

####R12 = f_R12(th, st, gr)

###################################################
###################################################
####
#### Construct K Matrix
#### following Method/Notation in Dietrich & Newsam (96)
####function f_K(th, st, gr)
####
####  st_ij = st.st_ij; n2 = size(st_ij)[1]
####  ## No of Stations
####
####  p = gr.p; q = gr.q
####
####  bar_rho = f_bar_rho(th, gr)
####
####  Lambda = sqrt(4*p*q)*Diagonal(real.(f_W(bar_rho, gr)))
####  #### THEORY CHECKPOINT:
####  #### All values in Lambda must be POSITIVE;
####  #### OTHERWISE, the METHOD FAILS!!!
####
####  bar_R21 = f_bar_R21(th, st, gr); bar_R12 = bar_R21'
####
####  WHbar_R12 = zeros(Complex, 4*p*q, n2)
####  for i in 1:n2
####
####    WHbar_R12[:,i] = f_WH(bar_R12[:,i], gr)
####
####  end
####
####  KH = Lambda^(-1/2)*WHbar_R12
####
####  K = KH' ## this gives Cnjugate Transpose
####
####  return(K)
####
####end
####
###################################################
### Example #######################################

#### K = f_K(th, st, gr)

##### Check!! #####################################

#### sort(eigen(bar_R11).values)
#### THEORY CHECKPOINT:
#### All eigenvalues of bar_R11 must be POSITIVE;
#### OTHERWISE, the METHOD FAILS!!!

#### A = bar_R21*inv(bar_R11)*(bar_R21)'

#### B = real.(K*K')

#### Check: Must Have A=B!!

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

#### bar_RR = [bar_R11 (bar_R21)'; bar_R21 R22]

#### THEORY CHECKPOINT:
#### All eigenvalues of bar_RR must be POSITIVE;
#### OTHERWISE, the METHOD FAILS!!!
#### sort(eigen(bar_RR).values)

#### Lambda = sqrt(4*p*q)*Diagonal(real.(f_W(bar_rho, gr)))

#### [sort(eigen(bar_R11).values) sort(eigen(Lambda).values)]
#### Check: The two colums must be the same !!

###################################################
###################################################
#### V is (n2)x1; result is (n1)x1
function f_A1_T(V, st_p, n1)

  n2 = length(V); A1_T_V = zeros(n1)

  for i in 1:n2

    A1_T_V[st_p[i]] = V[i]

  end

  return(A1_T_V)

end

###################################################
### Example #######################################

####V = rand(n2)

####A1_T_V = f_A1_T(V, st_p, n1)

###################################################
###################################################
#### U is (n1)x1; result is (n2)x1
function f_A1(U, st_p)

  n1 = length(U); n2 = length(st_p)

  A1_U = zeros(n2)

  for i in 1:n2

    A1_U[i] = U[st_p[i]]

  end

  return(A1_U)

end

###################################################
### Example #######################################

####U = rand(n1)

####A1_U = f_A1(U, st_p)

###################################################
###################################################
#### V is (n1)x1; result is (n1)x1
function f_Sigma(V, th, gr)

  p = gr.p; q = gr.q; n1 = (p+1)*(q+1); bar_n1 = 4*p*q

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

  bar_V = zeros(bar_n1)

  for j in 0:q

    bar_V[(j*2*p+1):(j*2*p+p+1)] = V[(j*(p+1)+1):((j+1)*(p+1))]

  end

  bar_rho = f_bar_rho(th, gr)

  Lambda = sqrt(bar_n1)*Diagonal(real.(f_W(bar_rho, gr)))
  #### THEORY CHECKPOINT:
  #### All values in Lambda must be POSITIVE;
  #### OTHERWISE, the METHOD FAILS!!!

  bar_Sigma_V = f_W(Lambda*f_WH(bar_V, gr),gr)

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

  Proj_Index = fill(0, n1)

  for j in 0:q

    Proj_Index[(j*(p+1)+1):((j+1)*(p+1))] = collect((j*2*p+1):(j*2*p+p+1))

  end

  Sigma_V = bar_Sigma_V[Proj_Index]

  return(real.(Sigma_V))

end

###################################################
### Example #######################################

####V = rand(n1); Sigma_V = f_Sigma(V, th, gr)

####maximum(abs.(R11*V-Sigma_V)) #### difference must be 0

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

####################################################################
####################################################################
### ALL FUNCTIONS SO FAR ARE STAND-ALONE ONES;
### THE ONE THAT FOLLOW NEED INTEGRATION WITH TUOMAS/MOSE CODE;
### THE POSITION ASSUMED IN THE CODE THAT FOLLOWS IS THAT WE ARE
### AT TIME t-1, AND HAVE N PARTICLES, EACH OF DIMENSION 3x(n1),
### REPRESENTING THE 3-DIMENSIONAL VECTOR FIELD ON A GRID OF n1 POSITIONS,
### AT TIME t-1.
### THE FUNCTIONS THE FOLLOW APPLY A SINGLE PARTICLE FILTER STEP.
### IN PARTICULAR, THE CODE THAT FOLLOWS DOES THE FOLLOWING:
### - IT SIMULATES THE OPTIMAL PROPOSAL FOR THE HEIGHT FIELD AT TIME t.
### - IT *DOES NOT* SIMULATE THE PROPOSAL FOR THE TWO VELOCITY FIELDS.
### - FOR THE LATTER PURPOSE THE CODE MUST BE COMBINED WITH THE CODE OF
###   OF TUOMAS/MOSE
###################################################################
####################################################################

#### Calculation of the mean for the optimal proposal for the
#### Height Field at time t; it is assumed that INTEGRATION with TUOMAS/MOSE code
#### it is assumed that INTEGRATION with TUOMAS/MOSE code will provide the
#### (N)x(n1) field FH_t which is the push-forward of all N particles at time t-1
#### under the PDE dynamics; FH_t then considers only the (N)x(n1) matrix with
#### the values of the height field. Y_obs_t are the (n2)x1 data at time t.
#### THE RESULT IS A MATRIX (N)x(n1)
function Calculate_Mean(FH_t, th, st, Yobs_t, Sobs, gr)

  ######### Offline Calculations #########

  p = gr.p; q = gr.q; n1 = (p+1)*(q+1)

  st_ij = st.st_ij; st_p = st.st_p

  n2 = size(st_ij)[1]; N = size(FH_t)[1]
  ## No of Stations + No of Particles

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

  R22 = f_R22(th, st, Sobs, gr)

  invR22 = inv(R22)

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

  mu20 = Sobs^(-2)*(R22 - Diagonal(fill(Sobs^2,n2)))
  #### (n2)x(n2)

  mu21 = f_Sigma(f_A1_T(invR22*(mu20*Yobs_t), st_p, n1), th, gr)
  #### (n1)x1

  mu22 = f_Sigma((Sobs^(-2))*f_A1_T(Yobs_t, st_p, n1), th, gr)
  #### (n1)x1

  mu2 = mu22 - mu21
  #### (n1)x1
  ##########################

  mu1 = zeros(N, n1)

  mu  = zeros(N, n1)

  ##### the required mean is different for each of the N particles
  for i in 1:N

    mu10 = invR22*f_A1(FH_t[i,:],st_p)
    #### (n2)x1

    mu11 = f_Sigma(f_A1_T(mu10, st_p, n1), th, gr)
    #### (n1)x1

    mu1[i,:] = FH_t[i,:] - mu11
    #### (n1)x1

    mu[i,:] = mu1[i,:] + mu2
    #### (n1)x1

  end

  return(mu)

end

###################################################
### Example #######################################

####N = 10; FH_t = rand(N, n1)

####Calculate_Mean(FH_t, th, st, Yobs_t, Sobs, gr)

###################################################
###################################################
#
# INPUT:
#### FH_t: an (N)x(n1) matrix where N is the number of particles
#### and n1 the size of the grid on the 2D-domain.
#### The i-th row of FH_t (referring to the i-th particle)
#### is the n1-dimensional Height field calculated after solving the PDE,
#### from time instance (t-1) to time t, using the i-th particle at time t-1 as initial condition.
#### th: Parameters in Covariance Function
#### st: Information about Stations
#### Yobs_t: n2-dimensional vector corresponding to the observations from the n2 stations
#### Sobs: Standard Deviation of Observations at Stations
#### gr: Parameters Specifying Regular Grid
# OUTPUT:
#### The numbers of particles N must be EVEN!!
function Sample_Height_Proposal(FH_t, th, st, Yobs_t, Sobs, gr)

  ######### Offline Calculations: #########
  ######### i.e. will be the same for 0->1, 1->2, t-1->t, etc. #########

  st_ij = st.st_ij; st_p = st.st_p

  n2 = size(st_ij)[1]; N = size(FH_t)[1]
  ## No of Stations + No of Particles

  p = gr.p; q = gr.q; P = gr.P; Q = gr.Q

  bar_n1 = 4*p*q; n1 = (p+1)*(q+1)

  Dx = P/p; Dy = Q/q

  bar_rho = f_bar_rho(th, gr)

  Lambda = sqrt(bar_n1)*Diagonal(real.(f_W(bar_rho, gr)))
  #### THEORY CHECKPOINT:
  #### All values in Lambda must be POSITIVE;
  #### OTHERWISE, the METHOD FAILS!!!

  bar_R21 = f_bar_R21(th, st, gr); bar_R12 = bar_R21'

  WHbar_R12 = zeros(Complex, bar_n1, n2)

  for i in 1:n2

    WHbar_R12[:,i] = f_WH(bar_R12[:,i], gr)

  end

  KH = Lambda^(-1/2)*WHbar_R12; K = KH' ## this gives Conjugate Transpose

  R22 = f_R22(th, st, Sobs, gr)

  L = cholesky(real.(R22-K*KH)).L
  #### THEORY CHECKPOINT:
  #### R22-K*KH must have POSITIVE EIGENVALUES,
  #### i.e. all(eigen(R22-K*KH).values)>0,
  #### OTHERWISE, the METHOD FAILS!!!

  R12 = f_R12(th, st, gr);

  R12_invR22 = R12*inv(R22)

  ##muZ1_C_Yobs = R12_invR22*Yobs_t

  Proj_Index = fill(0, n1)

  for j in 0:q

    Proj_Index[(j*(p+1)+1):((j+1)*(p+1))] = collect((j*2*p+1):(j*2*p+p+1))

  end

  ####

  Samples = zeros(N,n1)
  N0 = Int(N/2) ########## N must be even!

  ########################################
  ######### Online Calculations: ##########
  ######### here we need info from the data at time t, or/and a separate ##########
  ######### calculation for each particle. ##########
  means = Calculate_Mean(FH_t, th, st, Yobs_t, Sobs, gr)
  rng = Random.MersenneTwister(123)

  e1 = Vector{ComplexF64}(undef, bar_n1)
  e2 = Vector{ComplexF64}(undef, n2)

  for i in 0:(N0-1)

    @. e1 = complex(randn(rng), randn(rng))
    @. e2 = complex(randn(rng), randn(rng))

    bar_z1 = f_W(Lambda^(1/2)*e1, gr)
    z2 = (K*e1) + (L*e2)

    Rbar_z = real.([bar_z1; z2])
    Ibar_z = imag.([bar_z1; z2])

    Proj1 = [ Rbar_z[Proj_Index]; Rbar_z[(bar_n1+1):end] ]

    Proj2 = [ Ibar_z[Proj_Index]; Ibar_z[(bar_n1+1):end] ]

    ##Z1_C_Yobs_1 = muZ1_C_Yobs + Proj1[1:n1] - R12_invR22*Proj1[(n1+1):end]
    Z11 = Proj1[1:n1] - R12_invR22*Proj1[(n1+1):end]

    ##Z1_C_Yobs_2 = muZ1_C_Yobs + Proj2[1:n1] - R12_invR22*Proj2[(n1+1):end]
    Z12 = Proj2[1:n1] - R12_invR22*Proj2[(n1+1):end]

    ##Samples[2*i+1,:] = means[2*i+1,:] + Z1_C_Yobs_1
    Samples[2*i+1,:] = means[2*i+1,:] + Z11

    ##Samples[2*i+2,:] = means[2*i+2,:] + Z1_C_Yobs_2
    Samples[2*i+2,:] = means[2*i+2,:] + Z12

  end

  return(Samples)

end

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

###################################################
#### Example ######################################

N = 10; FH_t = rand(N, n1)

Ss = Sample_Height_Proposal(FH_t, th, st, Yobs_t, Sobs, gr)

####

Ss[1,st_p] - Yobs_t #### all the differences should be close to 0 if Sobs<<1

Ss[2,st_p] - Yobs_t #### all the differences should be close to 0 if Sobs<<1

Ss[3,st_p] - Yobs_t #### all the differences should be close to 0 if Sobs<<1

Ss[4,st_p] - Yobs_t #### all the differences should be close to 0 if Sobs<<1

####

Fs = zeros(N, q+1, p+1)

for i in 1:N

  Fs[i,:,:] = reshape(Ss[i,:], (p+1,q+1))'

end

i = 2

#heatmap(1:size(Fs[i,:,:],2),
#           1:size(Fs[i,:,:],1), Fs[i,:,:],
#           c=cgrad([:blue, :white,:red, :yellow]),
#           xlabel="x-axis", ylabel="y-axis",
#           title="")

#####!!!!!! Need to Check if I Combine in the Correct Way
#####!!!!!! the algorithm for the Simulation of the Conditional Field
#####!!!!!! and the Use of the Mean.

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


# This includes interpolation methods

export Interp1D, Interp2D, HInterp1D, HInterp2D, interpolate, project

const PointVector{Dim} = AbstractVector{<:AbstractPoint{Dim, T}} where {T}
const Tessellation = Union{<:LineString, <:PyObject}

abstract type AbstractInterp end
abstract type AbstractCurveInterp   <: AbstractInterp end       # One dimensional interpolations
abstract type AbstractSurfaceInterp <: AbstractInterp end       # Two dimensional interpolations

"""
    $TYPEDEF

One dimensional fractal interpolation. IFS is defined as
```math
    w_n(x, y) = \\begin{bmatrix}
        a_{11, n} & 0 \\
        a_{21, n} & a_{22, n}
    \\end{bmatrix}
    \\begin{bmatrix}
        b_{1, n} \\
        b_{2, n}
    \\end{bmatrix} \\quad n = 1, 2, \\ldots, N
```
Here ``a_{22, n}, n = 1, 2, \\ldots, N`` are the free variables. While constructing an `Interp1D`, the free varaibles can be
specfied as either a single real value or a vector of real values. If a single real value is specified, it assumed that
``a_{22} = a_{22, n}, n = 1, 2, \\ldots, N``.

# Fields

    $TYPEDFIELDS
"""
struct Interp1D{T<:Union{<:Real, <:AbstractVector{<:Real}}} <: AbstractCurveInterp
    """Free variables"""
    freevars::T
end

"""
    $TYPEDEF

Two dimensional fractal interpolation. IFS is defined as
```math
    w_n(x, y) = \\begin{bmatrix}
        a_{11, n}   & a_{12, n}     & 0         \\
        a_{21, n}   & a_{22, n}     & 0         \\
        a_{31, n}   & a_{32, n}     & a_{33, n} \\
    \\end{bmatrix}
    \\begin{bmatrix}
        b_{1, n} \\
        b_{2, n} \\
        b_{3, n} \\
    \\end{bmatrix} \\quad n = 1, 2, \\ldots, N
```
Here ``a_{33,n}, n = 1, 2, \\ldots, N` are the free variables. While constructing an `InterpdD`, the free variables can be
specfied as either a single real value or a vector of real values. If a single real value is specified, it assumed that
``a_{33, n}, n = 1, 2, \\ldots, N``.

# Fields

    $TYPEDFIELDS
"""
struct Interp2D{T<:Union{<:Real, <:AbstractVector{<:Real}}} <: AbstractSurfaceInterp
    """Free variables"""
    freevars::T
end

"""
    $TYPEDEF

One dimensional hidden fractal interpolation. IFS is defined as
```math
    w_n(x, y) = \\begin{bmatrix}
        a_{11, n}   & 0             & 0         \\
        a_{21, n}   & a_{22, n}     & a_{23, n} \\
        a_{31, n}   & a_{32, n}     & a_{33, n} \\
    \\end{bmatrix}
    \\begin{bmatrix}
        b_{1, n} \\
        b_{2, n} \\
        b_{3, n} \\
    \\end{bmatrix} \\quad n = 1, 2, \\ldots, N
```
Here
```math
    \\begin{bmatrix}
        a_{22, n} & a_{23, n} \\
        a_{32, n} & a_{33, n}
    \\end{bmatrix}
```
are the free variables. While constructing an `HInterp1D`, the free variables can be specfied as either a single real matrix
or a vector of real matrices. If a single real matrix is specified, it assumed that
```math
    \\begin{bmatrix}
        a_{22} & a_{23} \\
        a_{32} & a_{33}
    \\end{bmatrix}
    \\begin{bmatrix}
        a_{22, n} & a_{23, n} \\
        a_{32, n} & a_{33, n}
    \\end{bmatrix}
    \\quad \\forall n = 1, 2, \\ldots, N
```

# Fields

    $TYPEDFIELDS
"""
struct HInterp1D{T<:Union{<:AbstractMatrix, <:AbstractVector{<:AbstractMatrix}}} <: AbstractCurveInterp
    """Free variables"""
    freevars::T
end

"""
    $TYPEDEF

Two dimensional hidden fractal interpolation. IFS is defined as
```math
    w_n(x, y) = \\begin{bmatrix}
        a_{11, n}   & a_{11, n}     & 0             & 0         \\
        a_{21, n}   & a_{22, n}     & 0             & 0         \\
        a_{31, n}   & a_{32, n}     & a_{33, n}     & a_{34, n} \\
        a_{41, n}   & a_{42, n}     & a_{43, n}     & a_{44, n} \\
    \\end{bmatrix}
    \\begin{bmatrix}
        b_{1, n} \\
        b_{2, n} \\
        b_{3, n} \\
        b_{4, n} \\
    \\end{bmatrix} \\quad n = 1, 2, \\ldots, N
```
Here
```math
    \\begin{bmatrix}
        a_{33, n} & a_{34, n} \\
        a_{43, n} & a_{44, n}
    \\end{bmatrix}
```
are the free variables. While constructing an `HInterp2D`, the free variables can be specfied as either a single real matrix
or a vector of real matrices. If a single real matrix is specified, it assumed that
```math
    \\begin{bmatrix}
        a_{33} & a_{34} \\
        a_{43} & a_{44}
    \\end{bmatrix}
    \\begin{bmatrix}
        a_{33, n} & a_{34, n} \\
        a_{43, n} & a_{44, n}
    \\end{bmatrix}
    \\quad \\forall n = 1, 2, \\ldots, N
```

# Fields

    $TYPEDFIELDS
"""
struct HInterp2D{T<:Union{<:AbstractMatrix, <:AbstractVector{<:AbstractMatrix}}} <: AbstractSurfaceInterp
    """Free variables"""
    freevars::T
end

# Note: For plotting recipes, Interpolant should a subtype of Function
"""
    $TYPEDEF

Interpolation object.

# Fields

    $TYPEDFIELDS
"""
struct Interpolant{T1<:IFS, T2, T3<:AbstractInterp} <: Function
    """IFS """
    ifs::T1
    """Interpolation function"""
    itp::T2
    """Interpolation method"""
    method::T3
end

(interp::Interpolant)(x...) = interp.itp(x...)

"""
    $SIGNATURES

Returns an [`Interpolant`](@ref) by interpolating `pts` using `method`. `f0` is the initial function. `niter` is the number
of iterations during the construction of the interpolant. `method` can be

* [`Interp1D`](@ref)
* [`HInterp1D`](@ref)
* [`Interp2D`](@ref)
* [`HInterp2D`](@ref)

# Example
```julia
julia> xi, dx, xf = 0., 0.01, 10.;

julia> x = collect(xi : dx : xf);

julia> y = sin.(x);

julia> pts = collect.(zip(x, y));

julia> interp = integrate(pts, Interp1D(0.001));  # One dimensional interpolation.
```
"""
interpolate(pts::AbstractVector{<:AbstractVector{<:Real}}, method::AbstractInterp; f0 = getinitf(method), niter::Int = 10) =
    interpolate(map(pnt -> Point(pnt...), pts), method, f0=f0, niter=niter)

function interpolate(pts::PointVector, method::AbstractInterp; f0 = getinitf(method), niter::Int = 10)
    tess = tessellate(pts, method)
    transforms = gettransforms(pts, method)
    mappings = getmappings(transforms, method)
    itp = wrap(f0, tess, mappings, niter)[1]
    Interpolant(IFS(transforms), itp, method)
end

getinitf(::Interp1D)   = x -> 0.
getinitf(::HInterp1D)  = x -> [0., 0.]
getinitf(::Interp2D)   = (x, y) -> 0.
getinitf(::HInterp2D)  = (x, y) -> [0., 0.]

"""
    $SIGNATURES

Projects `pts` on a lower dimensional space by dropping the `drop` number of indexes.
"""
project(pts::PointVector, drop::Int=1)     = [Point(pnt[1 : end - drop]...) for pnt in pts]
project(pts::PointVector{2}, ::Interp1D)   = project(pts, 1)
project(pts::PointVector{3}, ::HInterp1D)  = project(pts, 2)
project(pts::PointVector{3}, ::Interp2D)   = project(pts, 1)
project(pts::PointVector{4}, ::HInterp2D)  = project(pts, 2)

# We use tessellation to locate the points. The fractal interpolant that is returned by `interpolate` function is a piecewise
# function. So to calculate the interpolant `interp` at a point `pnt`, we need to locate `pnt` to find the correct subdomain.
# We use tessellations to locate the points. See also `locate` function

tessellate(pts::PointVector{2}, method::Interp1D)  = LineString(project(pts, method))
tessellate(pts::PointVector{3}, method::HInterp1D) = LineString(project(pts, method))
tessellate(pts::PointVector{3}, method::Interp2D)  = spt.Delaunay(project(pts, method))
tessellate(pts::PointVector{4}, method::HInterp2D) = spt.Delaunay(project(pts, method))

# Because of finite precison arithmetic, the location of a valid point may fail. To overcome this, we can switch to arbitrary
# precision arithmetic (by using BigFloat) in the expense of increasing the computationonal complexity. However, not to
# increase the computationonal complexity directly, the strategy employed here is this:  before switching completely to
# arbitrary precision arithmetic, we first try to locate the point with finite precision (Float64 precision). If the point
# cannot be located, we double the precision and try to locate the point again. We do this until maximum allowed
# precison(specified to be 1024 bits as MAXPREC). If MAXPREC is reached and the point still cannot be found, the point
# location fails with an error message.

_locate(pnt::AbstractPoint{1, T}, tess::LineString) where {T} = findfirst(((p1, p2),) -> p1[1] ≤ pnt[1] ≤ p2[1], tess)
_locate(pnt::AbstractPoint{2, T}, tess::PyObject)   where {T} = tess.find_simplex(pnt)[1] + 1
function locate(pnt::AbstractPoint, tess::Tessellation)
    n = _locate(pnt, tess)
    n == 0 || return n
    pnt = doubleprecision(pnt)
    locate(pnt, tess)
end

doubleprecision(pnt::AbstractPoint{N, <:Real}) where {N} = Point(BigFloat.(pnt)...)
function doubleprecision(pnt::AbstractPoint{N, <:BigFloat}) where {N}
    prec = 2 * precision(BigFloat)
    prec ≤ MAXPREC ? setprecision(prec) : error("Exceeded maximum allowed precision $MAXPREC")
    Point(BigFloat.(pnt)...)
end

# IFS coefficients of the interpolant is found by using a linear algrebraic equation system using the boundary conditions.
# Each subtransformation in the transformations of a IFS maps a larger domain (Line in case of curve interpolation and
# Triangle in case of surface interpolation) to a smaller domain. These mappings maps the boundary points of the larger
# domains to boundary points of the smaller domains. `partition` function returns these smaller domains.

partition(pts::PointVector, method::AbstractCurveInterp) = LineString(pts)
function partition(pts::PointVector, method::AbstractSurfaceInterp, tess::Tessellation=tessellate(pts, method))
    trifaces = [TriangleFace(val...) for val in eachrow(tess.simplices .+ 1)]
    GeometryBasics.Mesh(pts, trifaces)
end

# In case of curve interpolations, the boundary is the line that connects the endpoints of the interpolation domain. In case
# of surface interpolations, the interpolation domain is triangle, than that triangle is the boundary. If the interpolation
# domain is ngon (such as, tetragon, pentagon, hexagon, etc.) the triangle that can be drawn inside the convex hull of the
# points and that has the largest area is returned.

getboundary(pts::PointVector, ::AbstractCurveInterp) = Line(pts[1], pts[end])
function getboundary(pts::PointVector, method::AbstractSurfaceInterp)
    # TODO: Add methods for nonconvex boundaries #42
    hull = spt.ConvexHull(collect(hcat(collect.(project(pts, method))...)') )
    if length(hull.vertices) == 3   # Triangle
        Triangle(Point.(pts[hull.vertices .+ 1])...)
    else    # Ngon (tetragon, pentagon, hexagon, etc...)
        polygon = Ngon(
            SVector{length(hull.vertices)}(Point.(pts[hull.vertices .+ 1]))
        )
        msh = GeometryBasics.mesh(polygon)
        idx = argmax([area(trig.points) for trig in msh])
        msh[idx]
    end
end

# `gettransforms` returns tranforms that map outer domains to smaller domains in the interpolation domain.
function gettransforms(pts::PointVector, method::AbstractInterp)
    parts = partition(pts, method)
    n = length(parts)
    freevars = typeof(method.freevars) <: AbstractVector ? method.freevars : fill(method.freevars, n)
    boundary = getboundary(pts, method)
    map(((domain, freevar),) -> _gettransform(boundary, domain, freevar), zip(parts, freevars))
end

# Interp1D
function _gettransform(outline::Line, inline::Line, freevar::Real)
    outmat = collect(hcat(coordinates(outline)...)')
    inmat = collect(hcat(coordinates(inline)...)')
    inmat[:, end] -= outmat[:, end] * freevar
    outmat[:, end] .= 1
    a11, b1, a21, b2 = outmat \ inmat
    A = [a11    0;
         a21    freevar]
    b = [b1, b2]
	Transformation(A, b)
end

# HInterp1D
function _gettransform(outline::Line, inline::Line, freevar::AbstractMatrix)
    outmat = collect(hcat(coordinates(outline)...)')
    inmat = collect(hcat(coordinates(inline)...)')
    inmat[:, 2 : 3] -= outmat[:, 2 : 3] * freevar'
    outmat = [outmat[:, 1] ones(2)]
    a11, b1, a21, b2, a31, b3 = outmat \ inmat
    A = [a11    0               0;
         a21    freevar[1,1]    freevar[1,2];
         a31    freevar[2,1]    freevar[2,2]
         ]
    b = [b1, b2, b3]
	Transformation(A, b)
end

# Interp2D
function _gettransform(outtrig::Triangle, intrig::Triangle, freevar::Real)
    outmat = collect(hcat(coordinates(outtrig)...)')
    inmat = collect(hcat(coordinates(intrig)...)')
    inmat[:, end] -= outmat[:, end] * freevar
    outmat[:, end] .= 1
    a11, a12, b1, a21, a22, b2, a31, a32, b3 = outmat \ inmat
    A = [a11    a12    0;
         a21    a22    0;
         a31    a32    freevar
         ]
    b = [b1, b2, b3]
	Transformation(A, b)
end

# HInterp2D
function _gettransform(outtrig::Triangle, intrig::Triangle, freevar::AbstractMatrix)
    outmat = collect(hcat(coordinates(outtrig)...)')
    inmat = collect(hcat(coordinates(intrig)...)')
    inmat[:, 3:4] -= outmat[:, 3:4] * freevar'
    outmat = [outmat[:, 1:2] ones(3)]
    a11, a12, b1, a21, a22, b2, a31, a32, b3, a41, a42, b4 = outmat \ inmat
    A = [a11    a12    0                0;
         a21    a22    0                0;
         a31    a32    freevar[1, 1]    freevar[1, 2];
         a41    a42    freevar[2, 1]    freevar[2, 2];
         ]
    b = [b1, b2, b3, b4]
	Transformation(A, b)
end

# Here transforms and mappings are different. The type of a transform is FractalTools.Transformation. A
# FractalTools.Transformation object basically consists of the matrix A and vector b of an affine transformation w(x) = A * x
# + b. However, a mapping is a subtransformation (defined as Ln: Ω ↦ Ωₙ and Fn: Ω × Rᵐ ↦ Rᵖ, where Ω is the interpolation
# domain and Ωₙ is the subdomain obtained by partitioning the interpolation domian) of an affine transformation in an IFS.
# So, we compute the mappings (i.e. subtransformations) from transforms and use directly the mappings to compute the
# interpolant.

getmappings(transforms, method) = map(transform -> _getmapping(transform, method), transforms)

function _getmapping(transform, method::Interp1D)
    (a11, a21, _, a22), (b1, b2) = transform.A, transform.b
    linv = x -> (x - b1) / a11
    F = (x, y) -> a21 * x + a22 * y + b2
    (linv, F)
end

function _getmapping(transform, method::HInterp1D)
    (a11, a21, a31, _, a22, a32, _, a23, a33), (b1, b2, b3) = transform.A, transform.b
    linv = x -> (x - b1) / a11
    F = (x, y, z) ->  [a21 a22 a23; a31 a32 a33] * [x, y, z] + [b2, b3]
    (linv, F)
end

function _getmapping(transform, method::Interp2D)
    (a11, a21, a31, a12, a22, a32, _, _, a33), (b1, b2, b3) = transform.A, transform.b
    linv = (x, y) -> [a11 a12; a21 a22] \ ([x, y] - [b1, b2])
    F = (x, y, z) ->  a31 * x + a32 * y + a33 * z + b3
    (linv, F)
end

function _getmapping(transform, method::HInterp2D)
    (a11, a21, a31, a41, a12, a22, a32, a42, _, _, a33, a34, _, _, a43, a44), (b1, b2, b3, b4) = transform.A, transform.b
    linv = (x, y) -> [a11 a12; a21 a22] \ ([x, y] - [b1, b2])
    F = (x, y, z, t) ->  [a31 a32 a33 a34; a41 a42 a43 a44] * [x, y, z, t] + [b3, b4]
    (linv, F)
end

# Main iteration of the initial function `f0` for `niter` iterations.
# TODO: Calculate `niter` with respect to a given ϵ
wrap(f0, tess::Tessellation, mappings::AbstractVector{<:Tuple{T, S}}, niter::Int) where {T, S} =
    ((f0, tess, mappings)) |> ∘((wrapper for i in 1 : niter)...)

function wrapper((f, tess, mappings))
    function fnext(x...)
        pnt = Point(x...)
        n = locate(pnt, tess)
        # n == 0 && error("Point $pnt cannot be found.")
        linv, F = mappings[n]
        val = linv(x...)
        F(val..., f(val...)...)
    end, tess, mappings
end


# Local Geary test of Spatial Autocorrelation
struct LocalGeary <: AbstractLocalSpatialAutocorrelation
    n::Int
    C::Vector{Float64}
    p::Vector{Float64}
    Cperms::Matrix{Float64}
    Cpermsmean::Vector{Float64}
    Cpermsstd::Vector{Float64}
    z::Vector{Float64}
    q::Vector{Symbol}
    categories::Symbol
end

"""
    localgeary(x, W)
Compute the Local Geary test of spatial autocorrelation.

# Optional Arguments
- `permutations=9999`: number of permutations for the randomization test.
- `rng=default_rng()`: random number generator for the randomization test.
- `corrected=true`: divide the scaling factor by ``n-1`` instead of ``n``.
- `categories=:positivenegative`: assing observations to positive or negative spatial autocorrelation, or in combination with the `:moran` scatterplot.
"""
function localgeary(x::AbstractVector{T} where T<:Real, W::SpatialWeights; permutations::Int = 9999,
    corrected::Bool = true, categories::Symbol = :positivenegative,
    rng::AbstractRNG = default_rng())::LocalGeary

    (categories == :positivenegative) || (categories == :moran) || throw(ArgumentError("`categories` must be :positivenegative or :moran"))

    n = length(x)
    z = x .- mean(x)

    m2::Float64 = sum(z.^2)
    if corrected
        m2 = m2 ./ (n - 1)
    else
        m2 = m2 ./ n
    end

    # Auxiliar function to calculate Local Moran
    function localgeary_calc(zi::Number, wi::AbstractVector, zneighi::AbstractVector)::Float64
        return (1 ./ m2)  .* sum(wi .* (zi .- zneighi).^2 )
    end

    # Local Geary
    C = zeros(n)
    for i in 1:n
        zi = z[i]
        wi = weights(W, i)
        zneighi = z[neighbors(W, i)]
        C[i] = localgeary_calc(zi, wi, zneighi)
    end

    # Conditional randomizatoin
    Cperms = crand_local(permutations, z, W, localgeary_calc, rng)

    larger = sum(Cperms .>= repeat(C, 1, permutations), dims = 2)
    low = (permutations .- larger) .< larger
    larger[low] .= permutations .- larger[low]
    p = (larger .+ 1) ./ (permutations + 1)
    p = vec(p)

    Cpermsstd = vec(std(Cperms, dims = 2, corrected = false))
    Cpermsmean = vec(mean(Cperms, dims = 2))
    zval = (C .- Cpermsmean)  ./ Cpermsstd
    zval = vec(zval)

    # Classification
    Cmean = mean(C)
    Wz = slag(W, z)

    q = Array{Symbol}(undef, n)
    if categories == :positivenegative
        for i in 1:n
            if zval[i] < 0
                q[i] = :P
            else
                q[i] = :N
            end
        end
    elseif categories == :moran
        for i in 1:n
            if zval[i] < 0
                if (z[i] > 0) & (Wz[i] > 0)
                    q[i] = :HH
                elseif (z[i] < 0) & (Wz[i] < 0)
                    q[i] = :LL
                else
                    q[i] = :OP
                end
            else
                q[i] = :NE
            end
        end
    end

    return LocalGeary(n, C, p, Cperms, Cpermsmean, Cpermsstd, zval, q, categories)

end

score(x::LocalGeary) = x.C;

scoreperms(x::LocalGeary) = x.Cperms;

mean(x::LocalGeary) = x.Cpermsmean;

std(x::LocalGeary) = x.Cpermsstd;

zscore(x::LocalGeary) = x.z;

pvalue(x::LocalGeary) = x.p;

assignments(x::LocalGeary) = x.q;

testname(::LocalGeary) = "Local Geary";

function labelsorder(x::LocalGeary)
    if x.categories == :positivenegative
        return [:ns, :P, :N]
    elseif x.categories == :moran
        return [:ns; :HH; :LL; :OP; :NE];
    end
end

function labelsnames(x::LocalGeary)
    if x.categories == :positivenegative
        return ["Not Significant"; "Positive"; "Negative"]
    elseif x.categories == :moran
        return ["Not Significant"; "High-High"; "Low-Low"; "Other Positive"; "Negative"]
    end
end

function labelcolor(::LocalGeary, x::String)

    # :positivenegative categories
    if x == "P"
        catcolor = :red
        alpha = 1.0
    elseif x == "N"
        catcolor =:blue
        alpha = 1.0
    # :moran categories
    elseif x == "HH"
        catcolor = "#B2182B" # (178, 24, 43)
        alpha = 1.0
    elseif x == "LL"
        catcolor = "#EF8A62" # (239, 138, 98)
        alpha = 1.0
    elseif x == "OP"
        catcolor = "#FDDBC7" # (253, 219, 199)
        alpha = 1.0
    elseif x == "NE"
        catcolor = "#67ADC7" # (103, 173, 199)
        alpha = 1.0
    # not significant
    elseif x == "ns"
        catcolor = :lightgrey
        alpha = 0.4
    end

    return (catcolor, alpha)
end


export skewProductFun, skewpoints, skewtransform!, iskewtransform!

#
# These are bivariate constructors that avoid sampling the diagonal y = x
# in a TensorSpace of two of the same spaces.
#

function skewProductFun{D1,D2}(f::Function,sp::TensorSpace{Tuple{Chebyshev{D1},Chebyshev{D2}}};tol=100eps())
    for logn = 4:10
        X = coefficients(skewProductFun(f,sp,2^logn,2^logn+1;tol=tol))
        if size(X,1)<2^logn && size(X,2)<2^logn+1
            return ProductFun(X,sp;tol=tol,chopping=true)
        end
    end
    warn("Maximum grid size of ("*string(2^11)*","*string(2^11+1)*") reached")
    skewProductFun(f,sp,2^11,2^11+1;tol=tol)
end

function skewProductFun{D1,D2}(f::Function,sp::TensorSpace{Tuple{Laurent{D1},Laurent{D2}}};tol=100eps())
    for logn = 4:10
        X = coefficients(skewProductFun(f,sp,2^logn,2^logn;tol=tol))
        if size(X,1)<2^logn && size(X,2)<2^logn
            return ProductFun(X,sp;tol=tol,chopping=true)
        end
    end
    warn("Maximum grid size of ("*string(2^11)*","*string(2^11)*") reached")
    skewProductFun(f,sp,2^11,2^11;tol=tol)
end

function skewProductFun(f::Function,S::TensorSpace,M::Integer,N::Integer;tol=100eps())
    xy = ApproxFun.checkpoints(S)
    T = promote_type(eltype(f(first(xy)...)),eltype(S))
    ptsx,ptsy=skewpoints(S,M,N)
    vals=T[f(ptsx[k,j],ptsy[k,j]) for k=1:size(ptsx,1), j=1:size(ptsx,2)]
    ProductFun(skewtransform!(S,vals),S;tol=tol,chopping=true)
end

function skewtransform!{T,D1,D2}(S::TensorSpace{Tuple{Chebyshev{D1},Chebyshev{D2}}},M::Matrix{T};kindx::Int=1,kindy::Int=2)
    n=size(M,1)

    planc=plan_chebyshevtransform(M[:,1];kind=kindx)
    for k=1:size(M,2)
        M[:,k]=chebyshevtransform(M[:,k],planc;kind=kindx)
    end

    planr=plan_chebyshevtransform(vec(M[1,:]);kind=kindy)
    for k=1:n
        M[k,:]=chebyshevtransform(vec(M[k,:]),planr;kind=kindy)
    end
    M
end

function iskewtransform!{T,D1,D2}(S::TensorSpace{Tuple{Chebyshev{D1},Chebyshev{D2}}},M::Matrix{T};kindx::Int=1,kindy::Int=2)
    n=size(M,1)

    planc=plan_ichebyshevtransform(M[:,1];kind=kindx)
    for k=1:size(M,2)
        M[:,k]=ichebyshevtransform(M[:,k],planc;kind=kindx)
    end

    planr=plan_ichebyshevtransform(vec(M[1,:]);kind=kindy)
    for k=1:n
        M[k,:]=ichebyshevtransform(vec(M[k,:]),planr;kind=kindy)
    end
    M
end

function skewtransform!{T,D1,D2}(S::TensorSpace{Tuple{Laurent{D1},Laurent{D2}}},M::Matrix{T})
    n=size(M,1)

    planc=plan_transform(S[1],M[:,1])
    for k=1:size(M,2)
        M[:,k]=transform(S[1],M[:,k],planc)
    end

    planr=plan_transform(S[2],vec(M[1,:]))
    for k=1:n
        M[k,:]=process!(transform(S[2],vec(M[k,:]),planr))
    end
    M
end

function iskewtransform!{T,D1,D2}(S::TensorSpace{Tuple{Laurent{D1},Laurent{D2}}},M::Matrix{T})
    n=size(M,1)

    planc=plan_itransform(S[1],M[:,1])
    for k=1:size(M,2)
        M[:,k]=itransform(S[1],M[:,k],planc)
    end

    planr=plan_itransform(S[2],vec(M[1,:]))
    for k=1:n
        M[k,:]=itransform(S[2],iprocess!(vec(M[k,:])),planr)
    end
    M
end

skewpoints(d::TensorSpace,n,m;kwds...) = skewpoints(d,n,m,1;kwds...),skewpoints(d,n,m,2;kwds...)

function skewpoints{D1,D2}(d::TensorSpace{Tuple{Chebyshev{D1},Chebyshev{D2}}},n,m,k;kindx::Int=1,kindy::Int=2)
    ptsx=fromcanonical(d[1],chebyshevpoints(n;kind=kindx))
    ptsy=fromcanonical(d[2],chebyshevpoints(m;kind=kindy))
    promote_type(eltype(ptsx),eltype(ptsy))[fromcanonical(d,x,y)[k] for x in ptsx, y in ptsy]
end

function skewpoints{D1,D2}(d::TensorSpace{Tuple{Laurent{D1},Laurent{D2}}},n,m,k)
    ptsx=fromcanonical(d[1],fourierpoints(n))
    ptsy=fromcanonical(d[2],fourierpoints(m)+π/m)
    promote_type(eltype(ptsx),eltype(ptsy))[fromcanonical(d,x,y)[k] for x in ptsx, y in ptsy]
end

function process!{T}(c::Vector{T})
    N = length(c)
    twiddle = exp(im*π/N)
    twid = one(T)
    for n=2:2:N-1
        twid *= twiddle
        c[n] *= twid
        c[n+1] /= twid
    end
    if mod(N,2) == 0
        twid *= twiddle
        c[N] *= twid
    end
    c
end

function iprocess!{T}(c::Vector{T})
    N = length(c)
    twiddle = exp(-im*π/N)
    twid = one(T)
    for n=2:2:N-1
        twid *= twiddle
        c[n] *= twid
        c[n+1] /= twid
    end
    if mod(N,2) == 0
        twid *= twiddle
        c[N] *= twid
    end
    c
end


using BenchmarkTools
using Printf
using Random
using LinearAlgebra

# Get the number of iterations from the command line.
n, = size(ARGS)
if n < 1
    println("Usage: belief.jl N")
    println("       ---> Please specify the number of iterations.")
    exit()
end

N = parse(Int, ARGS[1])

"""
    beliefpropagation(N)

Runs the belief propagation algorithm N times.
"""
function beliefpropagation(N)
    dim = 5000
    x = ones(dim)
    A = (randn(dim,dim) .+ 1.0)/2.0

    for i = 1:N
        x = log.(A*exp.(x));
        x .-= log.(sum(exp.(x)));
    end
    x
end

println("--------------------------")
println(@sprintf "Belief calculations: %d" N)
println("--------------------------")

@btime beliefpropagation(N)