#
# The l-family checker
# 
# The graph l_{2n+2} is a complete graph K_{2n} with two more nodes 
# each of them connected with half the nodes of the complete core. These new  
# nodes are also joint by an edge.
#
# The nodes are writen as 
#
#       (e,k,k',e')
#
# where e,e'\in\{0,1\} represents if the external nodes are mutants or not, 
# and k,k'\in{0,1,...,n} the number of mutants of each halves of the complete 
# core. Using the symmetries of the graph, it is possible to reduce the nodes 
# to those with (lexicographically)
# 
#       k >= k'                 e >= e'
#

# The parametre n is:
nn = 2



#
# Reduces the state s to a cannonical form
#
def reduce_state(s):
    if s[1]<s[2]:
        return (s[3],s[2],s[1],s[0])
    if (s[1]==s[2]) and (s[3]>s[0]):
        return (s[3],s[2],s[1],s[0])
    return s


#
# Compute the possible (reduced) states for a given n, it returns the list of
# states and the position of the states in the list 
#
def ComputeStates(n):
    # Compute the states
    states = []
    for k in [0..n]:
        for kp in [0..k]:
            if k != kp:
                states.extend([ (0,k,kp,0), (1,k,kp,0), (0,k,kp,1), (1,k,kp,1) ])
            else:
                states.extend([ (1,k,kp,1), (1,k,kp,0), (0,k,kp,0) ])
                
    istates = { v:k for k,v in enumerate(states)}
    
    s0 = states[0]
    sn = states[-1]
    states[0] = (0,0,0,0)
    states[-1] = (1,n,n,1)
    states[istates[(0,0,0,0)]] = s0
    states[istates[(1,n,n,1)]] = sn
    istates[s0] = istates[(0,0,0,0)]
    istates[sn] = istates[(1,n,n,1)]
    istates[(0,0,0,0)] = 0
    istates[(1,n,n,1)] = len(states)-1
    return states, istates


#
# Compute the matrix Q and the vector b for the list
#
def ComputeQb(r,n,states,istates):
    
    # Compute matrix
    Q = matrix(QQ, [[0]*len(states)]*len(states) )
    
    for s in states:
        m = sum(s)
        Q[ istates[s], istates[s] ] = r * m + 2 * n + 2 - m
    
    
    for s in states:
        # Leftmost
        if s[0] == 1:
            Q[istates[s], istates[(1,s[1],s[2],1)]] -= r/(n+1)
            if s[1] == n:
                Q[istates[s],istates[s]] -= r*n/(n+1)
            else:
                Q[istates[s],istates[(s[0],s[1]+1,s[2],s[3])]] -= r*(n-s[1])/(n+1)
                Q[istates[s],istates[s]] -= r*s[1]/(n+1)
        else:
            Q[istates[s], istates[(0,s[1],s[2],0)]] -= 1/(n+1)
    
            if s[1] == 0:
                Q[istates[s],istates[s]] -= n/(n+1)
            else:
                Q[istates[s],istates[reduce_state((s[0],s[1]-1,s[2],s[3]))]] -= s[1]/(n+1)
                Q[istates[s],istates[s]] -= (n-s[1])/(n+1)
    
        # Rightmost
        if s[3] == 1:
            Q[istates[s], istates[(1,s[1],s[2],1)]] -= r/(n+1)
            if s[2] == n:
                Q[istates[s],istates[s]] -= r*n/(n+1)
            else:
                Q[istates[s],istates[reduce_state((s[0],s[1],s[2]+1,s[3]))]] -= r*(n-s[2])/(n+1)
                Q[istates[s],istates[s]] -= r*s[2]/(n+1)
        else:
            Q[istates[s], istates[(0,s[1],s[2],0)]] -= 1/(n+1)
    
            if s[2] == 0:
                Q[istates[s],istates[s]] -= n/(n+1)
            else:
                Q[istates[s],istates[(s[0],s[1],s[2]-1,s[3])]] -= s[2]/(n+1)
                Q[istates[s],istates[s]] -= (n-s[2])/(n+1)
    
    
        # Core left
        if s[1] == n:
            Q[istates[s],istates[(1,n,s[2],s[3])]] -= r/2 #n*r*1/(2*n)
            Q[istates[s],istates[s]] -= r*(n-1)/2 #n*r * (n-1)/(2*n)
            if s[2] == n:
                Q[istates[s],istates[s]] -= r*n/2 #n*r * n/(2*n)
            else:
                Q[istates[s],istates[reduce_state((s[0],n,s[2]+1,s[3]))]] -= r*(n-s[2])/2 #n*r*(n-s[2])/(2*n)
                Q[istates[s],istates[s]] -= r*s[2]/2 # n*r*s[2]/(2*n)
        elif s[1] == 0:
            Q[istates[s],istates[(s[3],0,0,0)]] -= 1/2 #n/(2*n)
            Q[istates[s],istates[s]] -= (n-1)/2 + n/2 #n*(n-1)/(2*n) + n*n/2*n
        else:
            Q[istates[s],istates[(1,s[1],s[2],s[3])]] -= s[1]*r/(2*n)
            Q[istates[s],istates[reduce_state((0,s[1],s[2],s[3]))]] -= (n-s[1])/(2*n)
    
            Q[istates[s],istates[s]] -= s[1]*r*(s[1]-1)/(2*n) + (n-s[1])*(n-s[1]-1)/(2*n)
            Q[istates[s],istates[(s[0],s[1]+1,s[2],s[3])]] -= s[1]*r*(n-s[1])/(2*n)
            Q[istates[s],istates[reduce_state((s[0],s[1]-1,s[2],s[3]))]] -= (n-s[1])*s[1]/(2*n)
            
            Q[istates[s],istates[s]] -= r*s[1]*s[2]/(2*n) + (n-s[1])*(n-s[2])/(2*n)
            
            Q[istates[s],istates[reduce_state((s[0],s[1],s[2]+1,s[3]))]] -= r*s[1]*(n-s[2])/(2*n)
            
            if s[2] != 0:
                Q[istates[s],istates[(s[0],s[1],s[2]-1,s[3])]] -= (n-s[1])*s[2]/(2*n)
        
        # Core right
        if s[2] == n:
            Q[istates[s],istates[(1,n,n,s[0])]] -= r/2 #n*r*1/(2*n)
            Q[istates[s],istates[s]] -= r*(2*n-1)/2  # n*r*(2*n-1)/(2*n) 
        elif s[2] == 0:
            Q[istates[s],istates[(s[0],s[1],0,0)]] -= 1/2 #n/(2*n)
            
            if s[1] == 0:
                Q[istates[s],istates[s]] -= (2*n-1)/2 #n*(2*n-1)/(2*n)
            else:
                Q[istates[s],istates[s]] -= (n-1)/2 + (n-s[1])/2 #n*(n-1)/2*n + n*(n-s[1])/(2*n)
                Q[istates[s],istates[reduce_state((s[0],s[1]-1,0,s[3]))]] -= s[1]/2 #n*s[1]/(2*n)
        else:
            Q[istates[s],istates[reduce_state((s[0],s[1],s[2],1))]] -= s[2]*r/(2*n)
            Q[istates[s],istates[(s[0],s[1],s[2],0)]] -= (n-s[2])/(2*n)
    
            Q[istates[s],istates[s]] -= s[2]*r*(s[2]-1)/(2*n) + (n-s[2])*(n-s[2]-1)/(2*n)
            Q[istates[s],istates[s]] -= r*s[2]*s[1]/(2*n) + (n-s[2])*(n-s[1])/(2*n)
    
            Q[istates[s],istates[reduce_state((s[0],s[1],s[2]+1,s[3]))]] -= s[2]*r*(n-s[2])/(2*n)
            
            Q[istates[s],istates[(s[0],s[1],s[2]-1,s[3])]] -= (n-s[2])*s[2]/(2*n)
            
            if s[1] != n:
                Q[istates[s],istates[(s[0],s[1]+1,s[2],s[3])]] -= r*s[2]*(n-s[1])/(2*n)
            
            Q[istates[s],istates[reduce_state((s[0],s[1]-1,s[2],s[3]))]] -= (n-s[2])*s[1]/(2*n)
    
    b=-Q[1:len(states)-1,len(states)-1]
    Q=Q[1:len(states)-1,1:len(states)-1]
    return Q,b


#
# Computes the average fixation probability function
#
def Phi(Q,b,n,istates):
    sol = Q.solve_right(b)
    return sol[istates[(1,0,0,0)]-1,0] / (n + 1) + sol[istates[(0,1,0,0)]-1,0] * n / (n + 1) 


#
# "Entry point"
#

# Compute the states for the given size
states, istates = ComputeStates(nn)

# A first estimate of the degree of the rational function Phi
num_fac = len(states) - 2
fits=[1..num_fac+1]
fits.extend([ 1/f for f in fits[1:-1] ])

# Compute the fixation probabilities for the selected fitnesses
fps = {
    f : Phi(*ComputeQb(f, nn, states, istates), n=nn, istates=istates)
    for f in fits
}

# Now we are ready to compute the coefficients of the function. 
# Construct the matrix of the linear system describing the function \Phi
X = matrix(QQ, [[1]*len(fits)]*len(fits) )
X[:,num_fac-2] = [ [r] for r in fits ]
X[:,-1] = [ [ -fps[r] ] for r in fits ]
X[:,-2] = [ [ -r*fps[r] ] for r in fits ]

y = vector(QQ, [(fps[r]-1)*r**num_fac for r in fits ] )

for i in xrange(num_fac-3,-1,-1):
    X[:,i] = X[:,i+1].elementwise_product( X[:,num_fac-2] )
    X[:,num_fac + i] = X[:,num_fac+i+1].elementwise_product( X[:,num_fac-2] )

# Solve it
sol = X.solve_right(y)

# Create the field of fuctions. Then x is regarded as the variable of the 
# field Q(x). You can bypass this step to simplify you later the expressions 
# (that is regarding x as a symbolic variable with no special meaning)
F.<x> = Frac(QQ['x'])

# Construct the function
FB = (sum( sol[num_fac-1-i]* x**i for i in [0..num_fac-1] ) + x**num_fac) / (sum( sol[-(i+1)]* x**i for i in [0..num_fac-1] ) + x**num_fac)
pretty_print(FB)

# and finally the delta
delta = FB - x**(2*nn+1)/sum( x**i for i in range(2*nn+2) )
pretty_print(delta)

# take a look to the numerator and denominator
pretty_print(delta.denominator())

pretty_print(delta.numerator()/(x-1)) # up to (x-1)