using ArgParse
using SCS.SCSSolver

using Nemo
if VERSION >= v"0.6.0"
   import Nemo.Generic.perm
end

using PropertyT

using Groups

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

function matrix_emb(MM::MatSpace, g::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::WreathProductElem)(A::MatElem)
    G = matrix_emb(parent(A), g)
    inv_G = matrix_emb(parent(A), inv(g))
    return G*A*inv_G
end

function (p::perm)(A::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)
    return p*A*inv(p)
end

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

function E(i::Int, j::Int, M::MatSpace, val=one(M.base_ring))
    @assert i≠j
    m = one(M)
    m[i,j] = val
    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, X::Bool=false)
   indexing = [(i,j) for i in 1:n for j in 1:n if i≠j]
   G = MatrixSpace(ZZ, n, n)
   if X
      S = [E(i,j,G,v) for (i,j) in indexing for v in [1, 100]]
   else
      S = [E(i,j,G,v) for (i,j) in indexing for v in [1]]
   end
   S = vcat(S, [inv(x) for x in S])
   return G, unique(S)
end

function SL_generatingset(n::Int, p::Int, X::Bool=false)
    p == 0 && return SL_generatingset(n, X)
    (p > 1 && n > 1) || throw("Both n and p should be positive integers!")
    info("Size(SL($n,$p)) = $(SLsize(n,p))")
    F = ResidueRing(ZZ, p)
    G = 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, [inv(x) for x in S])
    return G, unique(S)
end

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

function cpuinfo_physicalcores()
    maxcore = -1
    for line in eachline("/proc/cpuinfo")
        if startswith(line, "core id")
            maxcore = max(maxcore, parse(Int, split(line, ':')[2]))
        end
    end
    maxcore < 0 && error("failure to read core ids from /proc/cpuinfo")
    return maxcore + 1
end

function parse_commandline()
    settings = ArgParseSettings()

    @add_arg_table settings begin
        "--tol"
            help = "set numerical tolerance for the SDP solver"
            arg_type = Float64
            default = 1e-14
        "--iterations"
            help = "set maximal number of iterations for the SDP solver (default: 20000)"
            arg_type = Int
            default = 200000
        "--upper-bound"
            help = "Set an upper bound for the spectral gap"
            arg_type = Float64
            default = Inf
        "--cpus"
            help = "Set number of cpus used by solver"
            arg_type = Int
            required = false
        "-N"
            help = "Consider elementary matrices EL(N)"
            arg_type = Int
            default = 2
        "-p"
            help = "Matrices over field of p-elements (p=0 => over ZZ)"
            arg_type = Int
            default = 0
        "--radius"
            help = "Radius of ball B_r(e,S) to find solution over"
            arg_type = Int
            default = 2
        "-X"
            help = "Consider EL(N, ZZ⟨X⟩)"
            action = :store_true
    end

    return parse_args(settings)
end

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

function main()
   parsed_args = parse_commandline()

   if parsed_args["cpus"] ≠ nothing
      if parsed_args["cpus"] > cpuinfo_physicalcores()
          warn("Number of specified cores exceeds the physical core cound. Performance will suffer.")
      end
      BLAS.set_num_threads(parsed_args["cpus"])
   end

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

   if p == 0
      if parsed_args["X"]
         dirname = "oSL$(N)Z⟨X⟩"
      else
         dirname = "oSL$(N)Z"
      end
   else
      dirname = "oSL$(N)_$p"
   end

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

   dirname = "$(dirname)_$(upper_bound)_r=$radius"

   logger = PropertyT.setup_logging(dirname)

   info(logger, "Group:       $dirname")
   info(logger, "Iterations:  $iterations")
   info(logger, "Precision:   $tol")
   info(logger, "Upper bound: $upper_bound")

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

   solver = SCSSolver(eps=tol, max_iters=iterations, linearsolver=SCS.Direct)

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

   PropertyT.check_property_T(sett)
end

main()