push!(LOAD_PATH, "./")

using Nemo
using Groups
using WreathProducts

using GroupRings
using PropertyT

import Nemo.elements

using JLD

include("Projections.jl")

###############################################################################
#
#  Iterator protocol for Nemo.FinField
#
###############################################################################

type FFEltsIter{T<:Nemo.FinField}
    all::Int
    field::T

    function FFEltsIter(F::T)
        return new(Int(characteristic(F)^degree(F)), F)
    end
end
FFEltsIter{T<:Nemo.FinField}(F::T) = FFEltsIter{T}(F)

import Base: start, next, done, eltype, length

Base.start(A::FFEltsIter) = (zero(A.field), 0)
Base.next(A::FFEltsIter, state) = next_ffelem(state...)
Base.done(A::FFEltsIter, state) = state[2] >= A.all
Base.eltype(::Type{FFEltsIter}) = elem_type(A.field)
Base.length(A::FFEltsIter) = A.all

function next_ffelem(f::Nemo.FinFieldElem, c::Int)
    if c == 0
        return (f, (f, 1))
    elseif c == 1
        f = one(parent(f))
        return (f, (f, 2))
    else
        f = gen(parent(f))*f
        return (f, (f, c+1))
    end
end

import Nemo.elements
elements(F::Nemo.FinField) = FFEltsIter(F)

###############################################################################
#
#  Action of Premutations on Nemo.MatElem
#
###############################################################################

function (p::Nemo.perm)(A::Nemo.MatElem)
    length(p.d) == A.r == A.c || throw("Can't act via $p on matrix of size ($(A.r), $(A.c))")
    R = parent(A)
    inv_p = inv(p)
    return R(Nemo.matrix_repr(p))*A*R(Nemo.matrix_repr(inv_p))
end


function orbit_decomposition(G::Nemo.Group, E::Vector, rdict=GroupRings.reverse_dict(E))

    elts = collect(elements(G))

    tovisit = trues(E);
    orbits = Vector{Set{Int}}()

    for i in 1:endof(E)
        if tovisit[i]
            orbit = Set{Int}()
            a = E[i]
            for g in elts
                idx = rdict[g(a)]
                tovisit[idx] = false
                push!(orbit,idx)
            end
            push!(orbits, orbit)
        end
    end
    return orbits
end

function matrix_repr(g::WreathProductElem, E, E_dict)
   rep_matrix = spzeros(Int, length(E), length(E))
   for (i,e) in enumerate(E)
      j = E_dict[g(e)]
      rep_matrix[i,j] = 1
   end
   return rep_matrix
end

function chars(G::PermutationGroup)
   permtype_unsorted(σ::Nemo.perm) = [length(c) for c in cycles(σ)]
   permtype(σ::Nemo.perm) = sort(permtype_unsorted(σ))

   χ_id(σ::Nemo.perm) = 1

   χ_sgn(σ::Nemo.perm) = (-1)^parity(σ)

   function χ_reg(σ::Nemo.perm)
      fixed_points = countnz([(x == y? 1 : 0) for (x,y) in enumerate(σ.d)])
      return fixed_points - 1
   end

   χ_regsgn(σ::Nemo.perm) = (-1)^parity(σ)*χ_reg(σ)

   function χ_regviaS3(σ::Nemo.perm)
      @assert parent(σ).n == 4
      t = permtype(σ)
         if t == [1,1,1,1]
            result = 2
         elseif t == [2,2]
            result = 2
         elseif t == [1,3]
            result = -1
         else
            result = 0
         end
      return result
   end

   chars = [χ_id, χ_sgn, χ_regviaS3, χ_reg, χ_regsgn]

   if G.n == 1
      return chars[1:1]
   elseif G.n == 2
      return chars[1:2]
   elseif G.n == 3
      return [chars[1:2]..., chars[4]]
   elseif G.n == 4
      return chars[1:5]
   else
      throw("Characters for $G unknown!")
   end
end

function epsilon(i, g::DirectProducts.DirectProductGroupElem)
    return reduce(*, 1, ((-1)^isone(g.elts[j]) for j in 1:i))
end

function central_projection(RG::GroupRing, char::Function, T::Type=Rational{Int})
    result = RG(T)
    for g in RG.basis
        result[g] = char(g)
    end
    return convert(T, char(RG.group())//Int(order(RG.group))*result)
end

function rankOne_projections(G::PermutationGroup, T::Type=Rational{Int})
   RG = GroupRing(G)
   projections = [central_projection(RG, χ, T) for χ in chars(G)]

   if G.n == 1 || G.n == 2
      return  projections
   elseif G.n == 3
      rankone_projs = [
         projections[1],
         projections[2],
         1//2*(one(RG) - RG(RG.group([2,1,3])))*projections[3]
         ]
      return rankone_projs
   elseif G.n == 4
      rankone_projs = [
         projections[1],
         projections[2],
         1//2*(one(RG) - RG(RG.group([2,1,3,4])))*projections[3],
         1//2*(one(RG) - RG(RG.group([2,1,3,4])))*projections[4],
         1//2*(one(RG) + RG(RG.group([2,1,3,4])))*projections[5]]
      return rankone_projs
   else
      throw("Rank-one projections for $G unknown!")
   end

end

function rankOne_projections(BN::WreathProducts.WreathProduct, T::Type=Rational{Int})

   N = BN.P.n
    # projections as elements of the group rings RSₙ
   SNprojs_nc = [rankOne_projections(PermutationGroup(i), T) for i in 1:N]

   # embedding into group ring of BN
   RBN = GroupRing(BN)
   RFFFF_projs = [central_projection(GroupRing(BN.N), g->epsilon(i,g), T)
        for i in 0:BN.P.n]
   Qs = [RBN(q, g -> BN(g)) for q in RFFFF_projs]

   function incl(k::Int, g::perm, WP::WreathProduct=BN)
      @assert length(g.d) + k <= WP.P.n
      arr = [1:k; g.d .+ k; (length(g.d)+k+1):WP.P.n]
      return WP(WP.P(arr))
   end

    all_projs=[Qs[1]*RBN(p, g-> incl(0,g)) for p in SNprojs_nc[N]]

   for i in 1:N-1
        Sk_first = [RBN(p, g->incl(0,g)) for p in SNprojs_nc[i]]
        Sk_last  = [RBN(p, g->incl(i,g)) for p in SNprojs_nc[N-i]]
        append!(all_projs, [Qs[i+1]*p1*p2
                for (p1,p2) in Base.product(Sk_first,Sk_last)])
   end

   append!(all_projs, [Qs[N+1]*RBN(p, g-> incl(0,g)) for p in SNprojs_nc[N]])

   return all_projs
end

function Cstar_repr(x::GroupRingElem, matrix_reps)
    res = zeros(matrix_reps[1])
    for i in 1:length(parent(x).basis)
        res += x.coeffs[i]*matrix_reps[i]
    end
    return res
end

function orthSVD(M::AbstractMatrix)
    M = full(M)
    fact = svdfact(M)
    sings = fact[:S]
    M_rank = sum(fact[:S] .> maximum(size(M))*eps(eltype(fact[:S])))
    Ufactor = fact[:U]
    return Ufactor[:,1:M_rank]
end

function Uπ_matrices(P_matrices; orth=orthSVD)
    U_p_matrices = Vector{Array{Float64,2}}()
    for (i,p_mat) in enumerate(P_matrices)
        U_p = orth(p_mat)
        push!(U_p_matrices, U_p)
    end
    return U_p_matrices
end

function compute_orbit_data{T<:GroupElem}(logger, name::String, G::Group, S::Vector{T}; radius=2)
   isdir(name) || mkdir(name)

   info(logger, "Generating ball of radius 4")
   @time E4, sizes = Groups.generate_balls(S, G(), radius=2*radius);
   info(logger, "Reverse dict")
   @time E_dict = GroupRings.reverse_dict(E4)

   info(logger, "Product matrix")
   @time pm = GroupRings.create_pm(E4, E_dict, sizes[radius], twisted=true)
   RG = GroupRing(G, E4, E_dict, pm)
   Δ = PropertyT.splaplacian(RG, S)
   @assert GroupRings.augmentation(Δ) == 0
   save(joinpath(name, "delta.jld"), "Δ", Δ.coeffs)
   save(joinpath(name, "pm.jld"), "pm", pm)

   N = length(G.objectGroup.gens)
   info(logger, "WreathProduct")
   @time BN = WreathProduct(Nemo.FiniteField(2,1, "a")[1], PermutationGroup(N))

   info(logger, "Decomposing E into orbits of B$(N)")
   @time orbs = orbit_decomposition(BN, E4, E_dict)
   @assert sum(length(o) for o in orbs) == length(E4)
   save(joinpath(name, "orbits.jld"), "orbits", orbs)

   info(logger, "Action matrices")
   E2 = E4[1:sizes[radius]]
   @time BNactionE_mreps = [matrix_repr(g, E2, E_dict) for g in elements(BN)]

   info(logger, "Projections")
   @time BN_mps = rankOne_projections(BN);

   @time π_E_projections = [Cstar_repr(p, BNactionE_mreps) for p in BN_mps]

   info(logger, "Uπs...")
   @time Uπs = Uπ_matrices(π_E_projections);

   multiplicities = [size(U,2) for U in Uπs];
   info(logger, "multiplicities = $multiplicities")
   dimensions = [Int(p[BN()]*Int(order(BN))) for p in BN_mps];
   info(logger, "dimensions = $dimensions")
   @assert dot(multiplicities, dimensions) == sizes[radius]

   save(joinpath(name, "U_pis.jld"), "Uπs", Uπs, "dims", dimensions)
   return 0
end