# Sage plethysm code to accompany abacus histories papers
# Focus is on checking equations, identities, understanding, etc.
# Greg Warrington: Aug 4, 2020
# 

import sys
from sage.combinat.subset import Subsets

# Not sure if this is the recommended approach, but seems to be working
Sym = SymmetricFunctions(FractionField(QQ['q','t','z']))
(q,t,z) = Sym.base_ring().gens()

# Name important bases
s = Sym.schur()
m = Sym.monomial()
p = Sym.powersum()
h = Sym.homogeneous()
e = Sym.elementary()
Ht = Sym.macdonald().Ht()
H = Sym.macdonald().H()
HLP = Sym.hall_littlewood().P()
HLQ = Sym.hall_littlewood().Q()

######################################################################################
# The Bernstein creation operators - S_alpha - as implemented in Sage
def Sbop_sage(n, f):
    """ `n`-th Bernstein creation operator applied to `f`
        This code is from the sage math docs
    """
    res = f.parent().zero()
    if not f:
        return res
    max_degree = max(sum(m) for m, c in f)
    for i in range(max_degree + 1):
        if n + i >= 0:
            res += (-1)**i * h[n + i] * f.skew_by(e[i])
    return res

def Salpha(alpha,f,opdef=Sbop_sage):
    """ Apply Sage's definition of Bernstein creation operator indexed by a Composition to f
    """
    ans = f
    for x in reversed(alpha):
        ans = opdef(x,ans)
    return s(ans)
    
# Salpha([2,4],p[0])

######################################################################################
# Top-level operators X_b
# The implementations for H/C/B are all based on the expansion of S_b
# in terms of e_c^{perp} and h_{b+c}
def S_b(b,f):
    """ This computes S_b using (eq 16) from Abacus Histories (AH) writeup
    """
    ans = 0*q**0
    # while S_b is defined for negative b, need to make sure we don't
    # try to access an h indexed by a partition with a negative part
    for c in range(max(0,-b),f.degree()+1):
        ans += (-1)**c*h[b+c]*f.skew_by(e[c])
    return s(ans)

def H_b(b,f):
    """ Computation of H_b using (eq 18) from AH writeup
        Relies on definition of S_b
    """
    ans = 0*q**0
    for c in range(f.degree()+1):
        ans += q**c*S_b(b+c,f.skew_by(h[c]))
    return ans

def C_b(b,f):
    """ Computation of C_b using (eq 20) from AH writeup
        Relies on definition of S_b
    """        
    ans = 0*q**0
    for c in range(f.degree()+1):
        ans += q**(-c)*S_b(b+c,f.skew_by(h[c]))
    return ans*(-q)**(1-b)

def B_b(b,f):
    """ Computation of B_b using equation from AH writeup at end of Sec 2
    """        
    ans = 0*q**0
    for d in range(f.degree()+1):
        tmpans = 0*q**0
        for c in range(max(0,-b-d),f.degree()+1):
            tmpans += (-1)**(-c)*e[b+d+c]*f.skew_by(e[d]).skew_by(h[c])
        ans += q**d * tmpans
    return s(ans)

#################################################################################
# Alternative formulas for the various operators (ones for S_b and C_b dependent)
#################################################################################
def S_b_viaC(b,f):
    """ This computes S_b using Prop. 3.6 from HMZ
        Appears after (eq 20) in AH writeup 
        Uses definition of C_b
    """
    ans = 0*q**0
    for c in range(f.degree()+1):
        ans += (-q)**(b-1)*C_b(b+c,f.skew_by(e[c]))
    return ans

def C_b_viaS(b,f):
    """ This computes C_b using (eq 20) from AH
        Uses definition of S_b
    """
    ans = 0*q**0
    for c in range(f.degree()+1):
        ans += (-1/q)**(b-1)*q**(-c)*S_b(b+c,f.skew_by(h[c]))
    return ans
    
def B_b_omega(b,f):
    """ Computation of B_b using (eq 21) from AH writeup
        # cf. conjugation of H by omega; (3.3) HMZ
    """
    return H_b(b,f.omega()).omega()

####################################################################
# Plethystic definitions of the operators
####################################################################
def _sb_psub(nu):
    """ replace p_nu with plethysm appearing in definition of S_b
    This plethystic substitution appears as P[X-1/z] in (eq 15) of AH
    and is the special case of p_nu[X-1/z]
    """
    ans = 1*q**0
    for k in nu:
        ans *= (p([k]) - z**(-k))
            
    return ans    

def _hb_psub(nu):
    """ replace p_nu with plethysm appearing in definition of H_b
    This plethystic substitution appears as P[X + (q-1)/z] in (eq 17) of AH
    and is the special case of p_nu[X + (q-1)/z]
    """
    ans = 1*q**0
    for k in nu:
        ans *= (p([k]) + q**k/z**k - z**(-k))
            
    return ans    

def _cb_psub(nu):
    """ replace p_nu with plethysm appearing in definition of C_b
    This plethystic substitution appears as P[X+(1/q-1)/z] in (eq 19) of AH
    and is the special case of p_nu[X + (1/q-1)/z]
    """
    ans = 1*q**0
    for k in nu:
        ans *= (p([k]) + q**(-k)/z**k - z**(-k))
            
    return ans    

def _gen_pleth_sub(P,subfn):
    """ Perform P[X-1/z] substitution of (eq 15) in AH for arbitrary P
    """
    ans = 0*q**0
    f = p(P)
    N = f.degree()
    for n in range(N+1): # function is not homogeneous
        for ptn in Partitions(n):
            coeff1 = f.coefficient(ptn)
            nf = subfn(ptn)
            ans += coeff1*nf
            
    return p(ans)

def _gen_pleth(b,f,subfn):
    """ compute S_b(f) (eq 15) or H_b(f) (eq 17) using plethystic formula found in AH
    """
    P = p(f)
    N = P.degree()
    
    # Perform the plethystic substitution
    ans = _gen_pleth_sub(P,subfn)
    
    # Multiply by sum_k h_k*z**k
    # Power of 1/z from substitutions will be no higher than degree(P)
    tmp = 0*q**0
    for i in range(N+b+1):
        tmp += z**i*h([i])
    
    # force Laurant poly in z to be poly by multiplying by z**N
    # for ease in coefficient extraction
    ans *= tmp*z**N
    nf = s(ans)
    # now extract coeff of z**(b+N)
    M = ans.degree()
    ans2 = 0*q**0
    for m in range(M+1):
        for nptn in Partitions(m):
            coeff = nf.coefficient(nptn)
            num = coeff.numerator()
            denom = coeff.denominator()
            ans2 += num.coefficient({z:b+N})*s(nptn)/denom
    return ans2        

########################################################################
# Plethystic formulas for operators
########################################################################
def S_b_pleth(b,f):
    """ compute S_b(f) using plethystic formula found in (eq 15) of AH
    """
    return _gen_pleth(b,f,_sb_psub)

def H_b_pleth(b,f):
    """ compute H_b(f) using plethystic formula found in (eq 17) of AH
    """
    return _gen_pleth(b,f,_hb_psub)

def C_b_pleth(b,f):
    """ compute C_b(f) using plethystic formula found in (eq 19) of AH
    """
    return (-1/q)**(b-1)*_gen_pleth(b,f,_cb_psub)

def B_b_pleth(b,f):
    """ compute B_b(f) using (eqs 19 and 21) of AH
    """
    return H_b_pleth(b,f.omega()).omega()

##########################################################
# compute X_alpha(1) using eqs (1)-(4) of abacus histories
##########################################################

def Salpha_one(alpha):
    """ compute S_alpha operator applied to 1
    """
    ans = p([])
    for k in list(reversed(alpha)):
        ans = S_b(k,ans)
    return s(ans)

def Calpha_one(alpha):
    """ compute C_alpha operator applied to 1
    """
    ans = p([])
    for k in list(reversed(alpha)):
        ans = C_b(k,ans)
    return s(ans)

def Halpha_one(alpha):
    """ compute H_alpha operator applied to 1
    """
    ans = p([])
    for k in list(reversed(alpha)):
        ans = H_b(k,ans)
    return s(ans)

# def Balpha_one(alpha,method='plethysm',RtoL=False):
#     """ Compute Balpha_one using (eq 3.6) in HMZ
#     
#     # WARNING: Look carefully at method being used
#     Only stated for partitions. Does it work in general?
#     RtoL True means that last (smallest) part is applied first as is
#        the case for H, C and S. This is convention in LW notes.
#     RtoL False is the default and is the convention in HMZ.
#     """
#     ans = (-q)**(len(alpha)-add(alpha))
#     if not RtoL then:
#         alpha = alpha[::-1]
#     return ans*(Calpha_one_pleth(alpha[::-1],True).omega())