''' 
    This file computes the Chern-Mather class for each Schubert variety and the
    Chern-Schwartz-MacPherson class for each Schubert cell in a simply laced
    cominuscule Grassmannian, and verifies strong posititivity of the CSM
    classes. It also computes the local Euler obstructions and compares them
    with the value of the Kazhdan-Lusztig polynomials evaluated at q=1. The
    program does not compute these values for largest Schubert variety, i.e.,
    the manifold itself. To include this computation, set the variable
    include_full_manifold to True. The lie_type and rank of the underlying group
    are supplied as the first and second command line argument. The cominuscule
    Grassmannian corresponds to a parabolic subgroup which contains all but one
    simple root. By default, the missing root is taken to be the last root. A
    third optional argument is used to override this behaviour. The output is
    written in a subdirectory named data.
'''

''' Refernces:
    [AM06] Aluffi, Mihalcea, "Chern-Schwartz-MacPherson classes for Schubert
    cells in flag manifolds", Compos. Math., 152(12):2603-2625, 2016.
    [MS20] Mihalcea, Singh, "Mather classes and conormal spaces of Schubert
    varieties in cominuscule spaces", arXiv:2006.04842.
'''

from collections import Counter
from copy import deepcopy

'''
A Counter object behaves like an array with two differences.
Any constant can be an index, rather than just integers.
The default value at a missing index is 0.
'''
c=Counter()
print(c)
print(c[10])
print(c['x'])
c['x'] += 1
c[3]=2
c['a']=5
print(c)

'''
We initialize the Grassmannian manifold. 
'''
include_full_manifold = True 

lie_type='A'
rank=5
d=2


'''
We check that the Grasmannian is of simply laced cominuscule type.
'''

assert 1 <= d <= rank
if lie_type not in ['A', 'D', 'E']:
    raise ValueError("Not a simply laced cominuscule space.")
if lie_type == 'E' and rank not in [6,7]:
    raise ValueError("Rank must equal 6 or 7 in type E.")
if lie_type == 'D' and rank <= 3:
    raise ValueError("Rank must be greater that 3 in type D.")
if lie_type in ['D', 'E'] and rank != d:
    raise ValueError("Not a simply laced cominuscule space.")

'''
We initialize the Weyl group.
'''
W = CoxeterGroup( lie_type + str(rank))
print(W)

'''
We create a dictionary alpha, such that alpha[i] corresponds to the i^th simple root.
'''
S = [W.simple_root_index(i) for i in W.index_set()]
alpha = [W.roots()[i] for i in S]
alpha = {i+1: root for i, root in enumerate(alpha)}
print(alpha)

C= CartanMatrix([lie_type, rank])

C

''' root_of_reflection is a dictionary. The keys are the reflections r_beta in
    W, and the values the corresponding positive root beta.
'''
root_of_reflection = {
    W.reflections()[root]: root\
        for root in W.roots()\
        if root > 0
}
root_of_reflection

def pairing(a,b):
    ''' Takes as input two roots a and b. Returns the bilinear product <a^v,b>,
        where a^v is the coroot dual to a.
    '''
    return a*C*b

def generateWP():
    ''' This function returns a list of elements in W^P sorted by length.
    '''
    new = {W.one()}
    by_len = []
    while new:
        by_len.append(new)
        new = set()
        for w in by_len[-1]:
            for sref in W.simple_reflections():
                v = sref * w
                if v.length()>w.length() and all(v*alpha[i]>0 for i in range(1, d)) and all(v*alpha[i]>0 for i in range(d+1, rank+1)):
                    new.add(v)
    retval=[]
    for sub in by_len:
        for w in sub:
            retval.append(w)
    return retval

WP = generateWP()

if not include_full_manifold:
    WP = WP[:-1] 

print([''.join([str(x) for x in w.reduced_word()]) for w in WP])

'''
WP is a list whose entries are the elements of W^P (minimal representatives of W/W_P).
We write a function idx which takes as input an element w of W^P and returns the index of w in the list WP.
The Mather and CSM classes will be encoded in a matrix where the i^th column corresponds to the class of the i^th element of WP.
The idx function, along with the csm and cMa functions (see below) allow a convenient way to view the CSM and Mather class of a particular element of WP. 
'''
idx_in_basis={}
for i, w in enumerate(WP):
    idx_in_basis[w]=i
    
def idx(w):
    return idx_in_basis[W.from_reduced_word(w)]

def pretty(vec):
    ''' This is a helper function to pretty print the homology class stored in
        vec.
    '''
    return ' + '.join(f'{vec[key]}*{key.reduced_word()}' for key in vec)

w=WP[1]
e= WP[0]
v=WP[2]

vec = Counter({w:4, e:1}) 
# vec corresponds to the Schubert class 1[X_e]+ 4[X_{s_2}]
# We can check the coefficients
print(f"{vec[e]=}")
print(f"{vec[w]=}")
print(f"{vec[v]=}")
print()
print(f"{vec=}")
print()

## We can also print vec in a pretty format
print("vec=", pretty(vec))

def lUpdate( root, vec):
    ''' The lUpdate funciton implements the Chevalley formula. It takes as input 
        a root and a homology class vec (stored as a python dictionary), and updates
        vec according to the rule vec = c(L_root) vec. Here c(L_root) is the 
        total Chern class of the line bundle corresponding to root. Note that vec is 
        updated in-place, unlike in tOp.
    '''
    c = deepcopy(vec)
    for w in c:
        for v in w.lower_cover_reflections():
            p = -pairing(root, root_of_reflection[v])
            vec[w*v] += p * c[w]

## We apply L_{alpha_1} to vec
s1 = W.from_reduced_word([1])
root = root_of_reflection[s1]
print("Before applying L_{alpha_1}, we have vec =", pretty(vec))
lUpdate(root, vec)
print("After applying L_{alpha_1}, we have vec =", pretty(vec))

def tOp( k, vec):
    ''' tOp implements the (cohomological) Demazure-Lusztig operator. 
        The variable vec is a homology class in G/B. The function applies the 
        operator T_k to vec and returns a new dictionary containing T_k(vec).
        See Thm 6.4 of [AM06] for more on the operator T_k, and its relation to
        CSM classes.
    '''
    res = Counter()
    for w, cof in vec.items():
        if w * alpha[k] < 0: 
            res[w] -= cof
            continue
        res[w] -= cof
        v = w * W.simple_reflection(k)
        res[v] += cof
        for u in v.lower_cover_reflections(): 
            beta = root_of_reflection[u]
            res[v*u] += cof * pairing( beta, alpha[k])
    return res

## We apply T_1 to vec
print("Before applying T_1, we have vec =", pretty(vec))
print("tOp(1,vec) = ", pretty(tOp(1,vec)))

'''
We compute the Mather classes of all elements of WP, and store it in cMaMatrix. 
'''
def compute_mather_classes():
    ''' This function computes the Mather class for each Schubert subvariety in
        G/P. We use Thm 1.1 of [MS20].
    '''
    Rn = [root for root in W.roots() if root < 0 and d-1 in root.support()]
    ''' Rn is the set of negative roots in the unipotent radical of the
        parabolic subgroup P.
    '''
    cMa=[]
    for w in WP:
        v = Counter([w])
        for beta in Rn:
            if w * beta > 0: 
                lUpdate(beta, v)
        cMa.append([])
        for u in WP:
            cMa[-1].append(v[u])
    return cMa

cMaMatrix = Matrix(compute_mather_classes()).transpose()
''' 
We have compute the expansion of the Mather class in the Schubert basis.
The i^th column is the Mather class of the i^th element in WP.
'''
print(cMaMatrix)
print("Basis : ", [w.reduced_word() for w in WP])

def compute_csm_classes():
    ''' This function computes the CSM class for each Schubert subvariety in
        G/P. Given w in W^P, we first compute the CSM class of X_B(w) using
        Thm 6.4 of [AM06], and then project this class to G/P.
    '''
    csm=[]
    for w in WP:
        v = Counter([W.one()])
        for k in w.reduced_word(): #T_{w^{-1}}*[X_id]
            v = tOp(k,v)
        csm.append([])
        for u in WP:
            csm[-1].append(v[u])
    # The following line verifies the strong positivity of the CSM class of X_w.
            assert u.bruhat_le(w) == (csm[-1][-1] != 0) 
    return csm


csmMatrix = Matrix(compute_csm_classes()).transpose()
''' 
We have compute the expansion of the CSM in the Schubert basis.
The i^th column is the CSM class of the i^th Schubert cell in WP.
'''
print(csmMatrix)
print("Basis : ", [w.reduced_word() for w in WP])

'''
We can now compute the Euler obstructions.
'''
eulerObsMatrix = csmMatrix.inverse() * cMaMatrix
print(eulerObsMatrix)
print("Basis : ", [w.reduced_word() for w in WP])

'''
We wish to generate a reduced word for the longest element of the Weyl subgroup W_P.
We identify the Dynkin diagram of W_P and use the the longest word.
For E6, the subdiagram of E6 obtained by removing alpha_6 has a different labelling than the standard labelling of D_5,
so we need to map the labels of D_5 to the corresponding root in E6.
'''

def longest_element_in_WP():
    if lie_type == 'E' and rank == 6:
        assert( d == rank)
        x = CoxeterGroup('D5').long_element().reduced_word()
        for i in range(len(x)):
            y = x[i]
            x[i] = 3 if y==2\
                else 4 if y==3\
                else 2 if y==4\
                else y
    elif lie_type == 'E' and rank == 7:
        assert( d == rank)
        x = CoxeterGroup('E6').long_element().reduced_word()
    elif lie_type == 'D':
        assert( d == rank)
        x = CoxeterGroup( 'A' + str( d - 1)).long_element().reduced_word()
    elif lie_type == 'A':
        x1 = x2 = []
        if d < rank:
            x2 = CoxeterGroup( 'A' + str(rank-d)).long_element().reduced_word()
            x2 = [ k + d for k in x2]
        if d > 1:
            x1 = CoxeterGroup( 'A' + str( d - 1)).long_element().reduced_word()
        x = x1 + x2
    return x

'''
We use the Coxeter3 package to generate the KL polynomials,
then sum up the coefficients of the KL polynomial to evaluate the value of the polynomial at q=1.
'''
def compute_kl():
    W3 = CoxeterGroup( lie_type + str(rank), implementation='coxeter3')
    W3P = [W3.from_reduced_word(x.reduced_word()) for x in WP]
    wP = W3.from_reduced_word(longest_element_in_WP())
    WPmax = [x*wP for x in W3P]
    klp = []
    for u in WPmax:
        klp.append([])
        for v in WPmax:
            poly = W3.kazhdan_lusztig_polynomial( u, v)
            klp[-1].append( sum(poly.coefficients()))
    return Matrix(klp)

klp = compute_kl()
print(klp)

'''
We verify that the value of the Kazhdan-Lusztig polynomial, evaluated at q=1, equals the Euler obstructions.
'''
if (klp == eulerObsMatrix):
    print("The Euler obstructions equal the Kazhdan-Lusztig polynomials evaluated at q = 1")
else:
    print("The Euler obstructions do not equal the Kazhdan-Lusztig polynomials evaluated at q = 1")

compute_kl(),compute_csm_classes(),eulerObsMatrix

'''
We write helper functions to output the Mather and CSM class of a particular element w.
These functions uses the cMaMatrix and csmMatrix to read off the Mather/CSM class.
'''

def cMa(w):
    i = idx_in_basis[w]
    res = ''
    for j, u in enumerate(WP):
        if cMaMatrix[j][i]!=0:
            res = res + f'+{cMaMatrix[j][i]}{u.reduced_word()}'

    res = res[1:]
    return res

def csm(w):
    i = idx_in_basis[w]
    res = ''
    for j, u in enumerate(WP):
        if csmMatrix[j][i]!=0:
            res = res + f'+{csmMatrix[j][i]}{u.reduced_word()}'

    res = res[1:]
    return res

print(cMa(W.from_reduced_word([3,1,2])))
print(csm(W.from_reduced_word([3,1,2])))