include("Orb_AutFN.jl")

###############################################################################
#
#  Action of WreathProductElems on Nemo.MatElem
#
###############################################################################

function matrix_emb(MM::Nemo.MatSpace, g::WreathProducts.WreathProductElem)
    parent(g).P.n == MM.cols == MM.rows || throw("No natural embedding of $(parent(g)) in ")
    powers = [(elt == parent(elt)() ? 0: 1) for elt in g.n.elts]
    elt = diagm([(-1)^(elt == parent(elt)() ? 0: 1) for elt in g.n.elts])
    return MM(elt)*MM(Nemo.matrix_repr(g.p)')
end

function (g::WreathProducts.WreathProductElem)(A::Nemo.MatElem)
    G = matrix_emb(parent(A), g)
    inv_G = matrix_emb(parent(A), inv(g))
    return G*A*inv_G
end

###############################################################################
#
#  Generating set
#
###############################################################################

function E(i::Int, j::Int, M::Nemo.MatSpace)
    @assert i≠j
    m = one(M)
    m[i,j] = m[1,1]
    return m
end

function SLsize(n,p)
    result = BigInt(1)
    for k in 0:n-1
        result *= p^n - p^k
    end
    return div(result, p-1)
end

function SL_generatingset(n::Int)
    indexing = [(i,j) for i in 1:n for j in 1:n if i≠j]
    G = Nemo.MatrixSpace(Nemo.ZZ, n, n)
    S = [E(i,j,G) for (i,j) in indexing];
    S = vcat(S, [transpose(x) for x in S]);
    return G, unique(S)
end

function SL_generatingset(n::Int, p::Int)
    p == 0 && return SL_generatingset(n)
    (p > 1 && n > 1) || throw("Both n and p should be positive integers!")
    @info("Size(SL($n,$p)) = $(SLsize(n,p))")
    F = Nemo.ResidueRing(Nemo.ZZ, p)
    G = Nemo.MatrixSpace(F, n, n)
    indexing = [(i,j) for i in 1:n for j in 1:n if i≠j]
    S = [E(i, j, G) for (i,j) in indexing]
    S = vcat(S, [transpose(x) for x in S])
    return G, unique(S)
end

###############################################################################
#
#  Parsing command line
#
###############################################################################

function parse_commandline()
    settings = ArgParseSettings()

    @add_arg_table settings begin
        "--tol"
            help = "set numerical tolerance for the SDP solver (default: 1e-5)"
            arg_type = Float64
            default = 1e-5
        "--iterations"
            help = "set maximal number of iterations for the SDP solver (default: 20000)"
            arg_type = Int
            default = 20000
        "--upper-bound"
            help = "Set an upper bound for the spectral gap (default: Inf)"
            arg_type = Float64
            default = Inf
        "--cpus"
            help = "Set number of cpus used by solver (default: auto)"
            arg_type = Int
            required = false
        "-N"
            help = "Consider automorphisms of free group on N generators (default: N=2)"
            arg_type = Int
            default = 2
        "-p"
            help = "Matrices over filed of p-elements (default: p=0 => over ZZ)"
            arg_type = Int
            default = 0
        "--radius"
            help = "Find the decomposition over B_r(e,S)"
            arg_type = Int
            default = 2
    end

    return parse_args(settings)
end

###############################################################################
#
#  main
#
###############################################################################

function main()
   parsed_args = parse_commandline()

   N = parsed_args["N"]
   p = parsed_args["p"]

   if p == 0
      dirname = "SL$(N)Z"
   else
      dirname = "SL$(N)_$p"
   end

   radius = parsed_args["radius"]
   dirname = "$(dirname)_r=$radius"

   isdir(dirname) || mkdir(dirname)
   logger = PropertyT.setup_logging(dirname)

   G, S = SL_generatingset(N, p)
   info(logger, G)
   info(logger, "Symmetric generating set of size $(length(S))")
   info(logger, S)
   AutS = WreathProduct(Nemo.FiniteField(2,1, "a")[1], PermutationGroup(N))

   tol = parsed_args["tol"]
   iterations = parsed_args["iterations"]
   upper_bound = parsed_args["upper-bound"]

   solver = SCS.SCSSolver(eps=tol, max_iters=iterations, verbose=true, linearsolver=SCS.Indirect)

   sett = Settings(dirname, N, G, S, AutS, radius, solver, upper_bound, tol)

   orbit_check_propertyT(logger, sett)
end

main()