import cc.redberry.core.context.OutputFormat
import cc.redberry.groovy.Redberry

import static cc.redberry.core.indices.IndexType.*
import static cc.redberry.core.tensor.Tensors.*
import static cc.redberry.groovy.RedberryPhysics.*
import static cc.redberry.groovy.RedberryStatic.*

//******************************************//
//******** Compton scattering in QCD *******//
//******************************************//

use(Redberry) {
    //setting up matrix objects
    defineMatrices 'T_A', Matrix2.matrix,//unitary matrices
            'G_a', Matrix1.matrix, //gamma matrices
            'u[p_a]', Matrix1.vector, Matrix2.vector, //quark wave function
            'cu[p_a]', Matrix1.covector, Matrix2.covector, //its conjugation
            'V_iA', Matrix1.matrix, Matrix2.matrix, //quark-gluon vertex
            'D[p_m]', Matrix1.matrix //quark propagator

    //SU(N) symmetric constants
    setSymmetric('d_ABC')
    //SU(N) structure constants
    setAntiSymmetric('f_ABC')

    //quark vertex
    def V = 'V_mA = -I*g*G_m*T_A'.t
    //quark propagator
    def D = 'D[p_m] = -I*(m + p_m*G^m)/(m**2 - p_m*p^m)'.t
    //sum over gluon polarizations in axial gauge
    def gluonPolarization = ('P_mn[k_a] =' +
            ' -g_mn + 1/(k_a*n^a)*(k_m*n_n + k_n*n_m) - 1/(k_a*n^a)**2 * k_m*k_n').t
    //auxiliary vector is unit
    def n2 = 'n_a*n^a = 1'.t
    //gluon propagator
    def G = 'G_mnAB[k_a] = I*g_AB*P_mn[k_a]/(k_a*k^a)'.t
    G = gluonPolarization >> G
    //3-gluon vertex
    def V3 =
            ('V_{mnr}^{ABC}[k1_m, k2_m, k3_m] =' +
                    'g*f^{ABC}*(g_mn*(k2_r - k1_r) + g_nr*(k3_m - k2_m) + g_mr*(k1_n-k3_n))').t

    //diagram a)
    def Ma = 'cu[p2_m]*V_mA*e^m[k2_m]*D[k1_m+p1_m]*V_nB*e^n[k1_m]*u[p1_m]'.t
    //diagram b)
    def Mb = 'cu[p2_m]*V_mB*e^m[k1_m]*D[p1_m-k2_m]*V_nA*e^n[k2_m]*u[p1_m]'.t
    //diagram c) (with 3-gluon vertex)
    def Mc = ('cu[p2_m]*V_mC*u[p1_m]*G^mnCD[p1_a-p2_a]*' +
            'V_{nabDBA}[p1_a-p2_a, k1_m, -k2_m]*e^a[k1_m]*e^b[k2_m]').t
    //matrix element
    M = Ma + Mb + Mc
    //substitute vertices and propagators in matrix element
    M = (D & G & V & V3) >> M
    //mandelstam and mass shell substitutions
    def mandelstam = setMandelstam(
            ['p1_m': 'm', 'k1_m': '0', 'p2_m': 'm', 'k2_m': '0'])
    //expand and apply substitutions
    M = (ExpandAll & EliminateMetrics & mandelstam) >> M

    //complex conjugation of matrix element
    def MC = Conjugate >> M
    //exchange spinor momentums
    MC = 'u[p1_m]*cu[p2_m] = u[p2_m]*cu[p1_m]'.t >> MC
    //reorder gamma and SU(N) matrices
    MC = (Reverse[Matrix1] & Reverse[Matrix2]) >> MC

    //squared matrix element
    //M2 = 'M_AB*MC^AB'
    def M2 = M * ('{_A -> ^A, _B -> ^B}'.mapping >> MC)
    //expand and eliminate contractions with metrics and deltas
    M2 = ExpandAndEliminate >> M2
    //sum over photon polarizations
    M2 = ('e_m[k1_a]*e_n[k1_a] = P_mn[k1_a]'.t & gluonPolarization) >> M2
    M2 = ('e_m[k2_a]*e_n[k2_a] = P_mn[k2_a]'.t & gluonPolarization) >> M2
    //sum over electron polarizations
    M2 = 'u[p2_m]*cu[p2_m] = m + p2^m*G_m'.t >> M2
    M2 = 'u[p1_m]*cu[p1_m] = m + p1^m*G_m'.t >> M2

    //trace of gamma matrices
    M2 = DiracTrace['G_a'] >> M2
    //trace of unitary matrices
    M2 = UnitaryTrace['T_A'.t, 'f_ABC'.t, 'd_ABC'.t, '3'.t] >> M2
    //simplifications and substitutions
    M2 = (ExpandAndEliminate & n2 & mandelstam) >> M2
    //simplify combinations of unitary constants
    M2 = UnitarySimplify['T_A'.t, 'f_ABC'.t, 'd_ABC'.t, '3'.t] >> M2

    //momentum conservation
    M2 = 'k2_a = p1_a + k1_a - p2_a'.t >> M2
    // replace scalar contractions with auxiliary vector n_m
    // (like p1_m*n^m) with symbols (e.g. p1_m*n^m = p1n)
    M2 = ExpandAll >> M2
    M2 = 'u = 2*m**2-s-t'.t >> M2
    ['p1', 'p2', 'k1'].each {
        M2 = "${it}_a *n^a = ${it}n".t >> M2
    }
    //mandelstam u
    M2 = 'u = 2*m**2-s-t'.t >> M2

    //write the result to file in Mathematica input form
    new File('quark-gluon-scattering').delete()
    new File('quark-gluon-scattering') << M2.toString(OutputFormat.WolframMathematica)
}