"""Given simple Lie group, for each non-distinguished rho determines a
single Weyl group element for which the score (zeta(w)_Delta in the article)
is non-positive.

Example usage:
>> python3 non-pos-weyl-elements.py F4
>> python3 non-pos-weyl-elements.py E6 > E6.txt

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

from sage.all import *
import itertools
import time
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(pos_roots, rho, w):
  score = 0

  for p in gen_V(2, rho, pos_roots):
    score += root_abs(w.action(p))

  for p in gen_V(0, rho, pos_roots):
    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(pos_roots, rho, w):
  return is_pos(compute_score(pos_roots, rho, w))

rholist = {}

# Distinguished rho for G2
rholist["G2"] = [
  (0,2),
  (2,2)
]

# Distinguished rho for F4
rholist["F4"] = [
  (0,2,0,0),
  (0,2,0,2),
  (2,2,0,2),
  (2,2,2,2)
]

# Distinguished rho for E6
rholist["E6"] = [
  (2,0,0,2,0,2),
  (2,2,2,0,2,2),
  (2,2,2,2,2,2)
]

# Distinguished rho for E7
rholist["E7"] = [
  (0,0,0,2,0,0,2),
  (2,0,0,2,0,0,2),
  (2,0,0,2,0,2,2),
  (2,2,2,0,2,0,2),
  (2,2,2,0,2,2,2),
  (2,2,2,2,2,2,2)
]

# Distinguished rho for E8
rholist["E8"] = [
  (0,0,0,0,2,0,0,0),
  (0,0,0,2,0,0,0,2),
  (0,0,0,2,0,0,2,0),
  (0,0,0,2,0,0,2,2),
  (2,0,0,2,0,0,2,0),
  (2,0,0,2,0,0,2,2),
  (2,0,0,2,0,2,0,2),
  (2,0,0,2,0,2,2,2),
  (2,2,2,0,2,0,2,2),
  (2,2,2,0,2,2,2,2),
  (2,2,2,2,2,2,2,2)
]

# Interpret command line arguments
# e.g.: python3 non-pos-weyl-elements.py F4
full_type = sys.argv[1]
n = int(full_type[1:])
sys_type = full_type[0]

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

print(R.dynkin_diagram())
print()

starttime = time.time()

counter = 0
total = 2**n

# Loop over every possible choice of rho
for rho in itertools.product([0, 2], repeat=n):
  # This ensures that we skip the distinguished rhos of the exceptional cases
  if full_type not in rholist.keys() or rho not in rholist[full_type]:
    pos = True

    for w in W:
      score = compute_score(P, rho, w)
      if not is_pos(score):
        pos = False
        break

    if not pos:
      # Complicated looking stuff to print out the Weyl group element and score in a pretty way
      tab_str = "   "
      w_mat_str = str(w.matrix()).replace("\n","\n" + tab_str + "  ")
      print(str(rho) + ":\n" + tab_str + "w=" + str(w) + "\n" + tab_str + "w=" + w_mat_str + "\n" + tab_str + "score=" + str(score) + "\n")
    else:
      print(str(rho) + " IS POSITIVE")

  # Progress output to stderr, in case stdout is piped to a file
  counter += 1
  sys.stderr.write(str(rho) + ": " + str(counter) + " out of " + str(total) + "\n")

endtime = time.time()

# sys.stderr.write("elapsed: " + str(endtime-starttime) + "\n")