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
