#!/usr/bin/env python
# coding: utf-8

# In[1]:


def x(i):
    return SR.var('x_{}'.format(i))


# In[2]:


def vcoding(lam,g,t):
    '''
    this function returns the V_{g,t} coding of a partition
    lam is the partition
    g is the value of the parameter depending on the type of Lie algebra considered
    t is the rank of the Lie algebra
    '''
    l=lam.length()
    tmp=[lam[i]-i-1  for i in range(l)]
    res = [abs(tmp[i]%g) for i in range(l)]
    #print(tmp)
    #print(set(res))
    #setres =set(res)
    cnt = 1
    while len(set(res))<t:
        tmp.append(-l-cnt)
        res.append((-l-cnt)%g)
        cnt+=1
    tmp_res=[]
    v=[]
    for i in range(len(tmp)):
        if res[i] not in tmp_res:
            v.append(tmp[i]+g)
            tmp_res.append(res[i])
        
    return v[:t]


# In[3]:


def g_value(ty,t):
    '''
    ty: type of the Lie algebra, can be a (for A^{(1)}), b (for B^{(1)}), bv for (for A_{odd}^{(2)}), c (for C^{(1)}),
    cv (for D_{t+1}^{(2)}), bc (for A_even^{(2)}) or d (for D^{(1)})
    this function computes the value of the parameter g for each type
    '''
    if ty=='a':
        g=t
    elif ty=='b':
        g=2*t-1
    elif ty=='c':
        g=2*t+2
    elif ty=='d':
        g=2*t-2
    elif ty=='bc':
        g=2*t+1
    elif ty=='cv':
        g=2*t
    elif ty=='bv':
        g=2*t
    return g


# In[4]:


def constant_coding(ty,t):
    '''
    constant coding is a function which takes the type ty of the Lie algebra and the rank t
    and return the constant c by which are shifted the elements of the V_{g,t}-coding in Section 4.2
    '''
    if ty=='a':
        c=0
    elif ty=='b':
        g=2*t-1
        c=g/2-1
    elif ty=='c':
        g=2*t+2
        c=g/2
    elif ty=='d':
        g=2*t-2
        c=g/2-1
    elif ty=='bc':
        g=2*t+1
        c=g/2
    elif ty=='bv':
        g=2*t
        c=g/2-1
        
    elif ty=='cv':
        g=2*t
        c=g/2-1/2
    return c


# In[5]:


def check_size(lam,ty,t):
    '''
    this function takes a partition lam , a type ty and the rank t of the Lie algebra.
    this function then computes the right-hand side of the first equation of all the Theorems of Section 4.2:
    it gives the size of lambda computed from the V_{g,t}-coding
    '''
    #value of the parameter g
    g=g_value(ty,t)
    #c is the constant to go from the V_{g,t}-coding to the vector (r_i) in the Theorems of Section 4.2
    c=constant_coding(ty,t)
    
    #ct is the constant we substract to the 1/g sum of (r_i)^2 (or in type A 1/(2g)sum of (r_i)^2)
    #to obtain the size of the corresponding partition
    if ty=='a':
        ct =(g-1)*(2*g-1)/12
    elif ty=='b':
        ct =(g+2)*(g//2+1)/12
    elif ty=='c':
        ct = (g/2-1)*(g-1)/12
    elif ty=='d':
        ct =(g+1)*(g//2+1)/12
    elif ty=='bc':
        ct =(g-2)*(g//2)/12
    elif ty=='cv':
        ct =(g*g-1)/24
    elif ty=='bv':
        ct =(g+1)*(g//2+1)/12
        
    v=vcoding(lam,g,t)
    
    tmp =sum([(v[i]-c)*(v[i]-c) for i in range(len(v))])
    if ty=='a':
        tmp=tmp/(2*g)-ct
    else:
        tmp=tmp/(g)-ct
    
    return tmp


# In[6]:


def Id(i,t):
    '''
    identity to set a default value function 
    '''
    return i


# In[7]:


def hook_functionty(lam,ty,t,x,x_hook=Id):
    '''
    this function returns both sides of Theorems of Section 4.2 as the first two outputs. Moreover it outputs the values of
    the signed statistic of the Lemma of Section 6.1. The last output is the product over boxes on the main diagonal which
    appears in the Lemmas from Section 6.1 when the formal variables are specialized. 
    
    
    lambda is a partition. It has to be chosen so that it matches the type of the affine Lie algebras. For instance
    it has to be a doubled distinct 6 cores if one looks at C_3^{(1)}
    
    ty: type of the Lie algebra, can be a (for A^{(1)}), b (for B^{(1)}), bv for (for A_{odd}^{(2)}), c (for C^{(1)}),
    cv (for D_{t+1}^{(2)}), bc (for A_even^{(2)}) or d (for D^{(1)})
    
    t is the rank
    
    x is the set of formal variables, note that it needs to be shifted because SR does not support negative index,
    it can therefore be specialized as a function. 
    
    WARNING: to avoid sagemath issues with handling variables with negative index, all of the indices of the variable are added +g.
    Therefore when specializing the set of variables as in Section 6.1, we have to shift our specializations by -g to recover 
    the computations from Section 6.1
   
    
    x_hook corresponds to the function necessary to compute the product on the main diagonal of Lambda 
    '''
    #Macdonald g value
    g=g_value(ty,t)
    
    #length of the partition
    l=lam.length()
    #length of the main diagonal
    d_l=lam.frobenius_rank()
    #sgn encodes the value of (-1)^{number of hook lengths <g such that Epsilon =1 (if type <>A)}
    sgn=1
    #sgn_diag encodes the value of (-1)^{number of hook lengths <g such that Epsilon =1 (if type <>A) that are on the main diagonal}
    sgn_diag=1
    #P is the product of hook lengths on the left-hand side of Theorems from Section 4.2
    P=1
    Q=1
    P_hookdiag=1
    for i in range(l):
        for j in range(lam[i]):
            h=lam.hook_length(i,j)
            #apart from the type A case, the products on the left-hand side of the theorems in Section 4.2 depend on the position
            #of a box with respect to the main diagonal
            if ty=='a':
                P=P*(x(h))/(x(h+g))*(x(h+2*g))/(x(h+g))
                if h<g:
                    Q=Q*(x(h))/(x(2*g-h))
                    sgn=-sgn
            else:
                if i<j:
                    P=P*(x(h))/(x(h+g))
                    if h<g:
                        sgn =-sgn
                        Q=Q*(x(h))/(x(2*g-h))
                elif i==j:
                    if ty=='d':
                        P=P*(x(h+2*g))/(x(h+g))
                    else:
                        P_hookdiag=P_hookdiag*(x_hook(h+g+g,ty))/x_hook(h-g+g,ty)
                        if h<g:
                            Q=Q*(x(h))/(x(2*g-h))
                            sgn=-sgn
                            sgn_diag = -sgn_diag
                        P=P*(x(h))/(x(h+g))
                elif i>j:
                    P=P*(x(h+2*g))/(x(h+g))
    vg=vcoding(lam,g,t)
    '''this part fills Q to be the right-hand side of Theorems from Section 4.2. 
    Since all of variables must be of positive index, the indices are shifted by g.
    Recall also that range(n) goes from 0 to n-1 which explains the (i+1) and the (j+1)
    '''
    for i in range(len(vg)):
        if ty=='c':
            Q=Q*x(vg[i]+g/2)/x(i+1+g)
        elif ty=='bv':
            Q=Q*x(vg[i]+g/2+1)/x(i+1+g)
        elif ty=='d':
            if x(vg[i]+g/2+1)==0:
                #here we just handle the exception r_i=0, then tau(2r_i)/tau(r_i)=1 in order to call 
                #hook_functionty for the applications to the q Nekrasov-Okounkov formula
                Q=Q
            else:
                Q=Q*x(2*vg[i]+2)/x(vg[i]+g/2+1)
            if i<len(vg)-1:
                Q=Q*x(i+1+g)/x(2*(i+1)+g)

        for j in range(i+1,len(vg)):
            if ty=='a':
                Q=Q*x(vg[i]-vg[j]+g)/x(j-i+g)
            elif ty in ['c','bc']:
                Q=Q*x(vg[i]-vg[j]+g)/x(j-i+g)*x(vg[i]+vg[j])/x(g-(i+1)-(j+1)+g)
            elif ty in ['d','b','bv']:
                Q=Q*x(vg[i]-vg[j]+g)/x(j-i+g)*x(vg[i]+vg[j]+2)/x(g+2-(i+1)-(j+1)+g)
            elif ty=='cv':
                Q=Q*x(vg[i]-vg[j]+g)/x(j-i+g)*x(vg[i]+vg[j]+1)/x(g+1-(i+1)-(j+1)+g)
    return [P,Q,sgn,sgn_diag,P_hookdiag]


# In[8]:


lambda_d=Partition([4,4,4,3,3])
testd=hook_functionty(lambda_d,'d',3,x)
print(testd[0]-testd[1])   


# In[9]:


lambda_bv=Partition([4,4,2,2,2])
testbv=hook_functionty(lambda_bv,'bv',10,x)
print('test')
print(testbv[0]-testbv[1])    


# In[10]:


lambda_d=Partition([4,4,4,3,3])
testd=hook_functionty(Partition([8,5,2,2,2,2,1,1,1]),'d',3,x)
print(testd[0]-testd[1])   


# In[11]:


def det_type(v,ty,x):
    '''
    v is the vector corresponding to the r_i
    ty stands for the type
    x is a set of formal variable
    this function returns the matrix whose determinant is either the numerator or the denominator of a schur function
    '''
    t=len(v)
    tmp=[]
    if ty=='a':
        for i in range(t):
            tmp.append([x(i)^(v[j]) for j in range(t)])
    elif ty=='d':
        for i in range(t):
            tmp.append([x(i)^(v[j])+x(i)^(-v[j]) for j in range(t)])
    else:
        for i in range(t):
            tmp.append([x(i)^(v[j])-x(i)^(-v[j]) for j in range(t)])
    return matrix(tmp)
        
        


# In[12]:


def hook_functiontychar(lam,t,ty,y,x):
    
    '''
    lam : partition
    
    t: rank of the Lie alebra
    
    ty: type of the Lie algebra, can be a (for A^{(1)}), b (for B^{(1)}), bv for (for A_{odd}^{(2)}), c (for C^{(1)}),
    cv (for D_{t+1}^{(2)}), bc (for A_even^{(2)}) or d (for D^{(1)})
    
    y : set of variables which are specialized as the function tau in the Lemmas of Section 6.1
    and in the checks performed below
    
    x : set of variables used to compute the characters. These are to be chosen in the applications
    as powers of the variable q to perform the checks of the Lemmas of Section 6.1
    
    this function returns the character 
    '''
    
    #c is the constant you have to substract from the V_g,t coding to obtain the r_i from Theorems of Section 4.2
    g=g_value(ty,t)
    c=constant_coding(ty,t)
    
    var('q')
    
    def x_hook(i,ty):
        '''
        auxillary function which computes the function appearing in the product over the boxes over Delta in Lemmas in Section 6.1
        
        ty is the type
        
        ct is the constant 1/2 when the terms of the product are quotients of terms of the form  1+/- q^{+/-g+h_s/2}
        and 1 if the terms are of the form  1+/- q^{+/-g+h_s}
        '''
        ct=1
        sgn=1
        if ty=='a':
            g=t
        elif ty=='b':
            g=2*t-1
            ct=1
            sgn =-1
        elif ty=='c':
            g=2*t+2
            ct=1/2
        elif ty=='d':
            g=2*t-2
            ct=1
        elif ty=='bc':
            g=2*t+1
            ct=1
            sgn =-1
        elif ty=='cv':
            g=2*t
            ct=1
            sgn =-1

        elif ty=='bv':
            g=2*t
            ct=1/2
        return 1+sgn*q^(ct*(i-g))
        
    def x_app(x,g):
        if ty=='a':
            g=t
        elif ty=='b':
            g=2*t-1
            c=g/2-1
        elif ty=='c':
            g=2*t+2
            c=g/2
        elif ty=='d':
            g=2*t-2
            c=g/2-1
        elif ty=='bc':
            g=2*t+1
            c=g/2
        elif ty=='bv':
            g=2*t
            c=g/2-1

        elif ty=='cv':
            g=2*t
            c=g/2-1/2

        def tmp_fun(i):
            return x(i-g)
        return tmp_fun
    
    l=lam.length()
    d_l=lam.frobenius_rank()
    vg=vcoding(lam,g,t)
    vg0=vcoding(Partition([]),g,t)
    r=[vg[i]-c for i in range(len(vg))]
    r0=[vg0[i]-c for i in range(len(vg0))]
    
    tmp=hook_functionty(lam,ty,t,x_app(y,ty),x_hook)
    print('function tau in 0 is equal to')
    print(x_app(y,ty)(0))
    print("the function appearing in the products over the main diagonal in Lemmas 6.1 evaluated in 0 is")
    print(x_hook(0,ty))
    #P is the product of hook lengths
    P=tmp[0]
    #sgn is the signature of the corresponding affine Grassmannian element
    sgn=tmp[2]
    #sgn_diag is the (-1)^{number of hooks on the main diagonal Delta whose length<g}
    sgn_diag=tmp[3]
    #P_hook is the correcting factor which is a product of boxes on the main diagonal
    P_hook=tmp[4]
    print('Product of theorems of Section 4.2 is')
    print(P)
    print('r coding is')
    print(r)
    #handling the case of type D character (eq 1.14)
    if ty=='d':
        if r[-1]==0:
            character=det_type(r,ty,x).det()/det_type(r0,ty,x).det()
        else:
            character=2*det_type(r,ty,x).det()/det_type(r0,ty,x).det()
    else:
        character=det_type(r,ty,x).det()/det_type(r0,ty,x).det()
    
    factor=1
    if ty=='c':
        factor=sgn*(1-x_app(y,ty)(g/2*d_l+g))
    elif ty=='b':
        factor=(1-x_app(y,ty)(-g/2*d_l+g))*sgn*sgn_diag
    elif ty=='bv':
        factor=(1-x_app(y,ty)(-g/2*d_l+g))*sgn
    elif ty=='cv':
        factor=(-1)^(d_l)
    elif ty=='bc':
        factor=sgn*sgn_diag*(1-x_app(y,ty)(d_l*g/2+g))
    elif ty=='d':
        factor=(1-x_app(y,ty)(-g*d_l+g))*sgn
    print('the power of (+/-u) appearing in the q Nekrasov Okounkov formula is when u is specialized to a power of q')
    print(factor)
    print('the product of hook length on the main diagonal specific to that type is ')
    print(P_hook)
    hook_prod=factor*P*P_hook
    return [character,hook_prod]
    


# In[13]:


def x_even(i):
    var('q')
    return 1-q^(2*(i))


# In[14]:


def x_clas(i):
    var('q')
    return 1-q^(i)


# In[15]:


def y_clas(i):
    var('q')
    return q^(i+1)


# In[16]:


def y_odd(i):
    var('q')
    return q^(2*i+1)


# In[17]:


def x_odd(i):
    var('q')
    return 1-q^(2*(i)+1)


# In[18]:


lambda_a=Partition([10,5])


# In[19]:


print(lambda_a.size())
print(check_size(lambda_a,'a',6))


# In[20]:


[test1a,test2a]=hook_functiontychar(lambda_a,6,'a',x_clas,y_clas)


# In[21]:


#testing if the expressions for Lemma 6.3 actually match
(test1a-test2a).canonicalize_radical().full_simplify()


# In[22]:


lambda_c=Partition([5,3,1,1])
print(lambda_c.size())
print(check_size(lambda_c,'c',2))


# In[23]:


[test1c,test2c]=hook_functiontychar(lambda_c,2,'c',x_clas,y_clas)


# In[24]:


#testing if the expressions for Lemma 6.3 actually match
(test1c-test2c).canonicalize_radical().full_simplify()


# In[25]:


lambda_bc=Partition([5,3,1,1])
print(lambda_bc.size())
print(check_size(lambda_bc,'bc',3))


# In[26]:


[test1bc,test2bc]=hook_functiontychar(lambda_bc,3,'bc',x_even,y_odd)


# In[27]:


#testing if the expressions for Lemma 6.10 actually match


# In[28]:


(test1bc-test2bc).canonicalize_radical().full_simplify()


# In[29]:


lambda_bv=Partition([4,1,1,1,1])
print(lambda_bv.size())
print(check_size(lambda_bv,'bv',2))


# In[30]:


[test1bv,test2bv]=hook_functiontychar(lambda_bv,2,'bv',x_clas,y_clas)


# In[31]:


#testing if the expressions for Lemma 6.6 actually match
(test1bv-test2bv).canonicalize_radical().full_simplify()


# In[32]:


lambda_d=Partition([6,5,4,4,4,2,1])


# In[33]:


print(lambda_d.size())
print(check_size(lambda_d,'d',3))


# In[34]:


[test1d,test2d]=hook_functiontychar(lambda_d,3,'d',x_even,y_odd)


# In[35]:


(test1d-test2d).canonicalize_radical().full_simplify()


# In[36]:


lambda_cv=Partition([4,4,2,2])
print(lambda_cv.size())
print(check_size(lambda_cv,'cv',2))


# In[37]:


[test1cv,test2cv]=hook_functiontychar(lambda_cv,2,'cv',x_even,y_odd)


# In[38]:


(test1cv-test2cv).canonicalize_radical().full_simplify()