from sage.all import *
from sage.geometry.polyhedron.constructor import *

from collections import defaultdict
from contextlib import nullcontext

class DistinguishedUnipotentOrbits:
  def __init__(self, sys_type, n, prefix=None):
    self.R = RootSystem([sys_type, n])
    self.L = self.R.root_lattice()

    self.sys_type = sys_type
    self.n = n

    self.W = self.L.weyl_group(prefix=prefix)
    self.P = self.L.positive_roots()
    self.Lw = self.R.weight_lattice()

    # Columns of this matrix will be coeff. of simple roots in fund. weights basis
    self._r2w = zero_matrix(QQ, self.n)

    for i in range(self.n):
      self._r2w[i] = self.Lw.simple_root(i+1).dense_coefficient_list()

    self._r2w = self._r2w.transpose()

    # Columns of this matrix will be coeff. of fund. weights in simple roots basis
    self._w2r = self._r2w.inverse()

    self._V = dict()
    self._ieqs = dict()
    self._Q = dict()
    self._Qw = dict()



  '''
  Utils
  '''

  def _eval_rho(self, rho, root):
    return sum(coeff * rho[delta-1] for (delta, coeff) in root)

  def _root_abs(self, root):
    if root.is_positive_root():
      return root
    else:
      return -root
      
  def compute_V(self, rho):
    if rho in self._V.keys():
      return self._V[rho]
    
    self._V[rho] = defaultdict(list)

    for r in self.P:
      self._V[rho][self._eval_rho(rho, r)].append(r)

    return self._V[rho]



  '''
  Score inequalities methods
  '''

  def _is_pos_score(self, score):
    if score == 0 or any(x <= 0 for x in score.dense_coefficient_list()):
      return False
    else:
      return True

  def compute_score(self, rho, w):
    score = 0
    V = self.compute_V(rho)

    for p in V[2]:
      score += self._root_abs(w.action(p))

    for p in V[0]:
      # The 2 * here comes from the fact that V[0] is only posiitve roots.
      score -= 2 * self._root_abs(w.action(p))

    return score
  
  def compute_all_scores(self, rho):
    return list(self.compute_score(rho, w) for w in self.W)
  
  def has_pos_score(self, rho, w):
    return self._is_pos_score(self.compute_score(rho, w))
  
  def has_all_pos_scores(self, rho):
    return all(self.has_pos_score(rho, w) for w in self.W)



  '''
  Cuspidal inequalities methods
  '''

  def compute_cusp_bounds(self, rho, w, upper_and_lower_bounds=False):
    upper_bounds = 0
    lower_bounds = 0
    V = self.compute_V(rho)

    for p in V[2]:
      wp = w.action(p)
      if wp.is_positive_root():
        upper_bounds += wp
      else:
        lower_bounds += wp

    # NOTE: this computation works because V0 only
    # contains the *positive* roots
    for p in V[0]:
      wp = w.action(p)
      abs_wp = self._root_abs(wp)
      upper_bounds -= abs_wp
      lower_bounds += abs_wp

    if upper_and_lower_bounds:
      return (upper_bounds, lower_bounds)
    else:
      return upper_bounds

  # Returns ieq_dict, of the form v.x <= b where ieq_dict[v] = b
  def compute_all_cusp_ineqs(self, rho):
    if rho in self._ieqs.keys():
      return self._ieqs[rho]
    
    self._ieqs[rho] = dict()

    for w in self.W:
      upper_bounds = self.compute_cusp_bounds(rho, w)

      # upper_bounds == 0 is an integer, so must be treated separately
      if upper_bounds == 0:
        upper_coeff = [0] * self.n
      else:
        upper_coeff = upper_bounds.dense_coefficient_list()
      
      for (row, upp) in zip(w.matrix(), upper_coeff):
        row_tuple = tuple(row)

        if row_tuple not in self._ieqs[rho]:
          self._ieqs[rho][row_tuple] = upp

    return self._ieqs[rho]
  
  # Returns cuspidal ineq polytope, in both simple root and fund. weight bases
  def compute_cusp_polytope(self, rho, weight_lattice_coord=False):
    if rho not in self._Q.keys():
      ieqs = self.compute_all_cusp_ineqs(rho)

      # For Polyhedron, tuple (a,b,c,d) corresponds to a + (b,c,d) @ x >= 0
      self._Q[rho] = Polyhedron(ieqs=list(
        (b,) + tuple(-v_i for v_i in v)
        for (v, b) in ieqs.items()))
    
    if weight_lattice_coord and rho not in self._Qw.keys():
      self._Qw[rho] = self._r2w * self._Q[rho]
    
    if weight_lattice_coord:
      return self._Qw[rho]
    else:
      return self._Q[rho]
  
  # Returns a list of tuples: (root lattice coord., weight lattice coord., is in interior)
  def compute_weight_lattice_points(self, rho):
    lattice_points = []

    Qw = self.compute_cusp_polytope(rho, weight_lattice_coord=True)

    for fw_x in Qw.integral_points():
      sr_x = self._w2r * fw_x
      # sr_coord = sum(sr_x[i] * self.L.simple_root(i+1) for i in range(self.n))
      # fw_coord = sum(fw_x[i] * self.Lw.fundamental_weight(i+1) for i in range(self.n))

      lattice_points.append((sr_x, fw_x, Qw.interior_contains(fw_x)))
      # lattice_points[fw_coord] = Qw.interior_contains(fw_x)

    return lattice_points
  
  def print_summary(self, rho_list, weight_lattice_coord=False, out_file=None):
    with open(out_file, "w") if out_file else nullcontext(sys.stdout) as f:
      for rho in rho_list:
        Q = self.compute_cusp_polytope(rho, weight_lattice_coord=weight_lattice_coord)
        lat = self.compute_weight_lattice_points(rho)
      
        f.write("rho: " + str(rho) + "\n")
        f.write("dim: " + str(Q.dim()) + "\n")
        f.write("facets: " + str(Q.n_facets()) + "\n")
        f.write("vertices: " + str(Q.n_vertices()) + "\n")
        f.write("lattice points:\n")

        for (v_sr,v_fw,in_interior) in lat:
          if weight_lattice_coord:
            v = v_fw
          else:
            v = v_sr

          if in_interior:
            f.write(str(v) + " <- interior\n")
          else:
            f.write(str(v) + "\n")

        f.write("\n")

  def center_is_lattice_point(self, rho_list, weight_lattice_coord=False, out_file=None):
    with open(out_file, "w") if out_file else nullcontext(sys.stdout) as f:
      for rho in rho_list:
        V = self.compute_V(rho)
        center = vector(sum(V[2])) / 2
        pt = self._r2w * center
        is_lattice = (pt.denominator() == 1)

        f.write("rho: " + str(rho) + "\n")
        if weight_lattice_coord:
          f.write("center: " + str(pt) + "\n")
        else:
          f.write("center: " + str(center) + "\n")
          
        f.write("lattice point: " + str(is_lattice) + "\n")
        f.write("\n")