import numpy as np
import time
import itertools
import sympy as sp
from sympy.matrices.exceptions import NonInvertibleMatrixError
import weyl_utils_numpy as wun
import weyl_utils as wu
import sys

# Command-line args from slurm
# E.g.: python3 check-rational-rho.py 0 16 E8 4
task_num = int(sys.argv[1])
task_total = int(sys.argv[2])
sys_type = sys.argv[3][0]
n = int(sys.argv[3][1:])

try:
  # Used to split combinations into two parts for more better parallelization
  # Specifically, don't have to generate all combinations for each task
  sub_k = int(sys.argv[4])
except:
  sub_k = n

# Generate combinations of n choose k, splitting the work by sub_k and total parts
def combinations_split_iter(n, k, sub_k, t, total):
  counter = 0
  for sub_comb in itertools.combinations(range(n - k + sub_k), sub_k):
    if counter % total == t:
      rest = range(sub_comb[-1]+1, n)
      for rest_comb in itertools.combinations(rest, k - sub_k):
        yield sub_comb + rest_comb
    counter += 1

file_prefix = './exceptional-data/'
DUO = wun.DistinguishedUnipotentOrbits(sys_type, n, file_prefix)
P = sp.Matrix(DUO.P)
twos = sp.Matrix([2] * n)

rho_list = []

print(sys_type + str(n))

print("--- Positive roots matrix ---")
for i in range(P.shape[0]):
    row_vector = list(P[i,:])
    print(f"{i}: {row_vector}")

print("--- Split combinations ---")

starttime = time.time()

for combo in combinations_split_iter(P.shape[0], n, sub_k, task_num, task_total):
  try:
    soln = P[combo,:].inv() @ twos

    if all(soln[i,0] >= 0 for i in range(soln.rows)):
      print("Combo:", combo)
      # print(P[combo,:])

      soln = tuple(soln)

      print("Soln:", soln)

      if not soln in wu.distinguished_rho_list(sys_type, n) \
          and not soln in rho_list:
        rho_list.append(soln)

      print("-")

  except NonInvertibleMatrixError:
    # print("Matrix is singular, no inverse.")
    pass

print("Total new rho found:", len(rho_list))
print("-----")

# Include the standard orbit as a double check
#rho_list.append((2,) * n)

DUO = wun.DistinguishedUnipotentOrbits(sys_type, n, file_prefix, rho_list=rho_list)

out = DUO.has_all_pos_scores()

distinguished_count = 0

for rho in out.keys():
  print(f"{rho}: all pos scores? {out[rho]}")
  if out[rho] and rho not in wu.distinguished_rho_list(sys_type, n):
    distinguished_count += 1

print("-----")

print(f"Total new distinguished unipotent orbits found for {sys_type}{n}: {distinguished_count}")

endtime = time.time()

print("-----")

print(f"Time taken for t={task_num}, total={task_total}: {endtime - starttime:.6f} seconds")