#!/usr/local64.hg/bin/sage -python

# Computes the Plancherel integral of the trace of the r-fold symmetric representation of SU(3,CC)

import sys
from sage.all import *

maxdim = 1
maxprime = 10

# specify polynomials
s = SymmetricFunctions(ZZ).schur()
representations = [s([r+1]).expand(3) for r in range(maxdim)]

# specify primes p
#primes = [2,3,5,7] # works
primes = list(primes(2, maxprime)) # doesn't work, out of memory.

# function setup
def e(x):
    return exp(2*pi*i*x)

# gamma(s) * gamma(-s) for a prime p and an angle theta
def gamma_sq(p,theta):
    return ( (1-1/p*e(-theta)) * (1-1/p*e(theta)) ) \
            / ( (1-e(-theta)) * (1-e(theta)) )

# |c(s)|^2 for all roots of sl3
def c_sq(p,a,b):
    return gamma_sq(p,a) * gamma_sq(p,b) * gamma_sq(p,-a-b)

# integrand = tr(repr'n) * factors for d mu(omega_s) -> ds
#           = tr(repr'n) * factors for ds->dtheta1,2 / |c(theta1,theta2)|^2
def integrand(p,poly,theta1,theta2):
    return  poly(e(theta1),e(theta2),e(-theta1-theta2))\
            * e(theta1+theta2)\
            / c_sq(p,theta1,theta2)

def calc_integrals(reps,primes):
    results = []
    theta1,theta2 = var('theta1,theta2')
    for rep in reps:
        print "Representation Polynomial: %s" % rep
        result_by_prime = []
        for prime in primes:
            print "Prime number %d: " % prime
            integral_by_theta1 = -4 *(pi**2) * integral(integrand(prime,rep,theta1,theta2),theta1,0,1)
            print "First integral: %s " % integral_by_theta1
            integral_by_theta2 = integral(integral_by_theta1, theta2,0,1)
            print "Second Integral: %s " % integral_by_theta2
            result_by_prime.append(integral_by_theta2)
        results.append(result_by_prime)
    return results

# annnd..... integrate!
results = calc_integrals(representations,primes)
o = open('calc_res_2.txt','w')
o.write(str(results))
o.close()