import numpy as np
import scipy.optimize as sciopt
from collections import defaultdict
import itertools
from numpy.lib.format import open_memmap
from weyl_utils import data_sizes, distinguished_rho_list

# from sage.all import *

class DistinguishedUnipotentOrbits:
  def __init__(self, sys_type, n, file_prefix, rho_list=None):
    self.sys_type = sys_type
    self.n = n

    self.dtype = "int8"

    if rho_list is None:
      # self.rho_list = distinguished_rho_list(sys_type, n)
      self.rho_list = list(itertools.product([0,2], repeat=self.n))
    else:
      self.rho_list = rho_list

    self._rho_array = np.array(self.rho_list)

    (_, num_P) = data_sizes(sys_type, n)
    
    fp_W_distinct = open_memmap(file_prefix + self.sys_type + str(self.n) + "-distinct.npy", mode='r')
    fp_P = np.memmap(file_prefix + self.sys_type + str(self.n) + "-pos-roots.dat", dtype=self.dtype, mode='r', shape=(num_P,n))

    self.W_distinct = np.array(fp_W_distinct)
    self.P = np.array(fp_P)

    self.WP_data = self.W_distinct @ self.P.T
    self.V_data = self.P @ self._rho_array.T

    self.V_map = dict()
    self.middle_x2 = dict()

    for i in range(len(self.rho_list)):
      rho = self.rho_list[i]

      self.V_map[rho] = defaultdict(lambda: np.zeros(self.P.shape[0], dtype=bool))
      
      for j in np.unique(self.V_data[:,i]):
        # NOTE: this only includes positive roots, even for V0
        self.V_map[rho][j] = (self.V_data[:,i] == j)

      self.middle_x2[rho] = np.sum(self.P[self.V_map[rho][2]], axis=0)



  '''
  Score inequalities methods
  '''

  def compute_all_scores(self):
    out = {rho: None for rho in self.rho_list}

    for rho in self.rho_list:
      # Equal to W_distinct @ V_i for i = 0,2
      WV2 = self.WP_data[:,self.V_map[rho][2]]
      WV0 = self.WP_data[:,self.V_map[rho][0]]

      s2 = np.sum(np.abs(WV2), axis=1)
      s0 = np.sum(np.abs(WV0), axis=1)
      total = s2 - 2 * s0

      out[rho] = total
    
    return out
  
  def has_all_pos_scores(self):
    out = {rho: None for rho in self.rho_list}

    scores = self.compute_all_scores()

    for rho in self.rho_list:
      out[rho] = np.all(scores[rho] > 0)

    return out



  '''
  Cuspidal inequalities methods
  '''

  def compute_all_cusp_bounds(self):
    upper_bounds = {rho: np.zeros(self.WP_data.shape[0], dtype=int) for rho in self.rho_list}

    for rho in self.rho_list:
      # Equal to W_distinct @ V_i for i = 0,2
      WV2 = self.WP_data[:,self.V_map[rho][2]]
      WV0 = self.WP_data[:,self.V_map[rho][0]]

      upper_bounds[rho] += np.sum(np.where(WV2 >= 0, WV2, 0), axis=1)

      # NOTE: this computation works because we only have the *positive* roots in V0
      upper_bounds[rho] -= np.sum(np.abs(WV0), axis=1)
      
    return upper_bounds
  
  # Output is of the form A @ x <= b
  def compute_all_cusp_ineqs(self):
    b = {rho: None for rho in self.rho_list}

    upper_bounds = self.compute_all_cusp_bounds()

    A = self.W_distinct

    for rho in self.rho_list:
      b[rho] = upper_bounds[rho]

    return (A, b)