#########################################################################################
# Code for doing some reality checks among all of the various definitions floating around
# Greg Warrington: Aug 4, 2020
#########################################################################################

def weak_compositions(n,k):
    if n < 0 or k  <0:
        return
    elif k == 0:
        if n == 0:
            yield []
        return
    elif k == 1:
        yield [n]
        return
    else:
        for i in range(n+1):
            for comp in weak_compositions(n-i,k-1):
                yield [i] + comp
                    
def strong_compositions(n,k):
    for c in weak_compositions(n,k):
        if 0 not in c:
            yield c
            
def compositions(n):
    ans = []
    for k in range(0,n+1):
        ans.extend(list(strong_compositions(n,k)))
    return list(ans)

######################################################################
# Check AH as a model for various operators
######################################################################
def _convert_lc(lc):
    """ convert lc to proper symmetric function in Schur basis
    """
    cdict = dict()
    for f in lc.d:
        ptn = list(filter(lambda x: x != 0, f['ah'].get_shape(r=-1)))
        if str(ptn) in cdict.keys():
            cdict[str(ptn)]['coeff'] += f['sgn']*q**f['pow']
        else:
            cdict[str(ptn)] = dict()
            cdict[str(ptn)]['coeff'] = f['sgn']*q**f['pow']
            cdict[str(ptn)]['ptn'] = ptn
    ans = 0*q**0
    for k in sorted(reversed(list(cdict.keys()))):
        ans += cdict[k]['coeff']*s(cdict[k]['ptn'])
    return ans

def check_Salpha(alpha,simplified=True):
    tmp = AH.S_alpha(alpha,simplified=simplified,verbose=False)
    
    ansh = _convert_lc(tmp)
    ansp = Salpha_one(alpha)
    if ansh != ansp:
        print("S_%s plethysic and AH not equal!" % (alpha))
        print("ansh %s" % (ansh))
        print("ansp %s" % (ansp))
        return False
    else:
        print("Success: %s" % (ansp))
        return True

def check_Halpha(alpha,simplified=True):
    tmp = AH.H_alpha(alpha,simplified=simplified,verbose=False)
    ansh = _convert_lc(tmp)
    ansp = Halpha_one(alpha)
    if ansh != ansp:
        print("H_%s plethysic and AH not equal!" % (alpha))
        print("ansh %s" % (ansh))
        print("ansp %s" % (ansp))
        return False
    else:
        print("Success: %s" % (ansp))
        return True

def check_Calpha(alpha,simplified=True):
    tmp = AH.C_alpha(alpha,simplified=simplified,verbose=False)
    ansh = _convert_lc(tmp)
    ansp = Calpha_one(alpha)
    if ansh != ansp:
        print("C_%s plethysic and AH not equal!" % (alpha))
        print("ansh %s" % (ansh))
        print("ansp %s" % (ansp))
        return False
    else:
        print("Success: %s" % (ansp))
        return True

####################################################################################
# Check that our AH formulas match the plethystic ones for compositions of n <= maxN
####################################################################################
def check_alphas(maxN):
    for n in range(maxN+1):
        for c in compositions(n):
            ans1 = check_Salpha(c)
            ans2 = check_Halpha(c)
            ans3 = check_Calpha(c)
            ans4 = check_Salpha(c,simplified=False)
            ans5 = check_Halpha(c,simplified=False)
            ans6 = check_Calpha(c,simplified=False)
            if not ans1 or not ans2 or not ans3 or not ans4 or not ans5 or not ans6:
                return False
    return True
        
# Reality checks of various relations

######################################################
# Check (eq 7) of AH: S_m S_n = - S_{n-1} S_{m+1}
######################################################
def check_S_comm_m_n(m,n):
    """ 
    There's no obvious upper limit of partition size to test against, using max(n+1,m+2)
    """
    maxN = max(n+1,m+2)
    for N in range(maxN+1):
        for mu in Partitions(N):
            lhs = s(S_b(m,S_b(n,s[mu])))
            rhs = -s(S_b(n-1,S_b(m+1,s[mu])))
            if lhs != rhs:
                print("S commutation relation failed for m=%d n=%d s[%s]" % (m,n,mu))
                print("LHS: %s" % (lhs))
                print("RHS: %s" % (rhs))
                return False
    return True

def check_S_comm(maxmn):
    """ Check S commutation relation for various values of m and n
    """
    for m in range(0,maxmn):
        for n in range(0,maxmn):
            check_S_comm_m_n(m,n)        
    print("S commutation check through indices of max size %s successful" % (maxmn))

#####################################################################################
# Check near (eq 21) of AH: q C_m C_n - C_{m+1} C_{n-1} = C_n C_m - q C_{n-1} C_{m+1}
#####################################################################################
def check_C_comm_m_n(m,n):
    """ 
    There's no obvious upper limit of partition size to test against, using max(n+1,m+2)
    """
    maxN = max(n+1,m+2)
    for N in range(maxN+1):
        for mu in Partitions(N):
            lhs = s(q * C_b(m,C_b(n,s[mu])) - C_b(m+1,C_b(n-1,s[mu])))
            rhs = s(C_b(n,C_b(m,s[mu])) - q * C_b(n-1,C_b(m+1,s[mu])))
            if lhs != rhs:
                print("C commutation relation failed for m=%d n=%d s[%s]" % (m,n,mu))
                print("LHS: %s" % (lhs))
                print("RHS: %s" % (rhs))
                return False
    return True

def check_C_comm(maxmn):
    """ 
    """
    for m in range(0,maxmn):
        for n in range(0,maxmn):
            check_C_comm_m_n(m,n)        
    print("C commutation check through indices of max size %s successful" % (maxmn))

####################################################
# Check that B_b and B_b_omega definitions are equal
####################################################

def check_B_b_defs(maxN):
    """ We need to be able to remove at least b boxes from s[mu] via skewing operators
    otherwise all formulas based on S_b will return 0
    """
    for N in range(1,maxN+1):
        for mu in Partitions(N):
            for b in range(-N-1,N+2):
                lhs = B_b(b,s[mu])
                rhs = B_b_omega(b,s[mu])
                if lhs != rhs:
                    print("B formulas unequal for b=%d s[%s]" % (b,mu))
                    print("LHS: %s" % (lhs))
                    print("RHS: %s" % (rhs))
                    return False
    print("B_b definitions match for |mu| <= %s" % (maxN))

##############################################################
# Check that S_b, S_b_viaC and S_b_pleth definitions are equal
##############################################################
def check_S_b_defs(maxN):
    """ We need to be able to remove at least b boxes from s[mu] via skewing operators
    otherwise all formulas based on S_b will return 0
    """
    for N in range(1,maxN+1):
        for mu in Partitions(N):
            for b in range(-N-1,N+2):
                lhs = S_b(b,s[mu])
                rhs = S_b_viaC(b,s[mu])
                whs = S_b_pleth(b,s[mu])
                if lhs != rhs or lhs != whs:
                    print("S formulas unequal for b=%d s[%s]" % (b,mu))
                    print("LHS: %s" % (lhs))
                    print("RHS: %s" % (rhs))
                    print("WHS: %s" % (whs))
                    return False
    print("S_b definitions match for |mu| <= %s" % (maxN))

####################################################
# Check that C_b and C_b_viaS definitions are equal
####################################################
def check_C_b_defs(maxN):
    """ We need to be able to remove at least b boxes from s[mu] via skewing operators
    otherwise all formulas based on S_b will return 0
    """
    for N in range(1,maxN+1):
        for mu in Partitions(N):
            for b in range(-N-1,N+2):
                lhs = C_b(b,s[mu])
                rhs = C_b_viaS(b,s[mu])
                if lhs != rhs:
                    print("S formulas unequal for b=%d s[%s]" % (b,mu))
                    print("LHS: %s" % (lhs))
                    print("RHS: %s" % (rhs))
                    return False
    print("C_b definitions match for |mu| <= %s" % (maxN))

####################################################
# Check that H_b and H_b_pleth definitions are equal
####################################################
def check_H_b_defs(maxN):
    """ We need to be able to remove at least b boxes from s[mu] via skewing operators
    otherwise all formulas based on S_b will return 0
    """
    for N in range(1,maxN+1):
        for mu in Partitions(N):
            for b in range(-N-1,N+2):
                lhs = H_b(b,s[mu])
                rhs = H_b_pleth(b,s[mu])
                if lhs != rhs:
                    print("S formulas unequal for b=%d s[%s]" % (b,mu))
                    print("LHS: %s" % (lhs))
                    print("RHS: %s" % (rhs))
                    return False
    print("S_b definitions match for |mu| <= %s" % (maxN))
                    

######################################################################
# Run the actual checks
######################################################################

# check_alphas(3)    # Check that X_alpha operator definitions match (pleth vs. abacus histories

# check_S_comm(4)    # check commutation relations
# check_C_comm(4)    

# check_B_b_defs(5)  # compare the various definitions we have for these operators
# check_S_b_defs(5)   
# check_H_b_defs(5)  
# check_C_b_defs(5)

# print(Halpha_one([1,1,1,1,1,1,1,1,1]))