"""Given simple Lie group and a rho vector, determines if the
score (zeta(w)_Delta in the article) of every Weyl group element is
positive or not.

Example usage:
>> python3 pos-check-parallel.py A4 "(2,0,2,0)"
>> python3 pos-check-parallel.py E8 "(2,2,2,2,2,2,2,2)"

Authors: Alexander Bertoloni Meli, Teruhisa Koshikawa, Jonathan Leake 
Date: August 2024
"""

from sage.all import *
from multiprocessing import Pool
import time
from ast import literal_eval
import sys

# Determine the weight of a given root based on rho
def eval_rho(rho, root):
  return sum(coeff * rho[delta-1] for (delta, coeff) in root)

# Generator function for V_k
def gen_V(k, rho, roots):
  for r in roots:
    if eval_rho(rho, r) == k:
      yield r

# Absolute value of a given root
def root_abs(root):
  return root.map_coefficients(abs)

def compute_score(V2, V0, w):
  score = 0

  for p in V2:
    score += root_abs(w.action(p))

  for p in V0:
    score -= 2 * root_abs(w.action(p))

  return score

def is_pos(score):
  if any(x <= 0 for x in score.dense_coefficient_list()):
    return False
  else:
    return True
  
def has_pos_score(V2, V0, w):
  return is_pos(compute_score(V2, V0, w))

starttime = time.time()

# Interpret command line arguments
# e.g.: python3 pos-check-parallel.py E8 "(2,2,2,2,2,2,2,2)"
sys_type = sys.argv[1][0]
n = int(sys.argv[1][1:])
rho = literal_eval(sys.argv[2])

# Initialize root system and Weyl group data
R = RootSystem([sys_type, n])
L = R.root_lattice()
W = L.weyl_group()
P = L.positive_roots()
Delta = L.simple_roots()

V2 = list(gen_V(2, rho, P))
V0 = list(gen_V(0, rho, P))

def pos_helper(w):
  return has_pos_score(V2, V0, w)

pos = True

# Parallelize the computation to make it run faster
with Pool(maxtasksperchild=1) as pool:
  # This means we apply the function pos_helper to every element of W
  # in a parallel way, and store the output in reslist.
  reslist = pool.imap_unordered(pos_helper, W, chunksize=500000)
  pool.close()

  for result in reslist:
    # If result is False, then it means some Weyl element has non-positive score
    if not result:
      pool.terminate()
      pos = False
      break

  pool.join()

if pos:
  print(str(rho) + " YES")
else:
  print(str(rho) + " NO")

endtime = time.time()

print("elapsed: " + str(endtime-starttime))