diff --git a/src/ipglasma.cpp b/src/ipglasma.cpp
index 97df1099db81a442f17547710ade0292b2f74b98..bcc31d2aabf172030ccfffe2398a34ef64cd7e3e 100644
--- a/src/ipglasma.cpp
+++ b/src/ipglasma.cpp
@@ -624,11 +624,14 @@ double IPGlasma::SumAmplitudes(int l)
 //TODO: check
 double IPGlasma::SumAmplitudes()
 {
-    double sum = 0.0;
     int L = xcoords.size();
+    double z = 0.5;
+    double Q = 0;
+
+ 
     cout << "# Summing with lattice size L = " << L << endl;
+    double sum = 0.0;
     WilsonLine quark;
-    //int thrownAway = 0;
     #pragma omp parallel for reduction(+:sum)
     for(int i =0; i < L; i += 1){
         for(int j =0; j < L; j += 1){
@@ -636,35 +639,22 @@ double IPGlasma::SumAmplitudes()
             int q1[2] = {i,j};
             quark = GetWilsonLine(WilsonLineCoordinate(i, j));
             double r1 = DistanceToOrigin(q1);
-            //#pragma omp parallel for reduction(+:sum) collapse(2)
     
             for(int k =0; k < L; k +=1){
                 for(int l =0; l < L; l +=1){
     
                     int q2[2] = {k,l};
                     double r2 = DistanceToOrigin(q2);
-                    if (r1 > 1. && r2 > 1.){
-                       // cout << "#!!Throwing away large dipole separation!" << endl;
-                        //cout << "#Quark coordinate indices " << q1[0] << "," << q1[1] << " and " << q2[0] << "," << q2[1] << endl;
-                        //cout << "#Quark coordinates " << X(q1[0]) << "," << Y(q1[1]) << " and " << X(q2[0]) << "," << Y(q2[1]) << endl;
-                        //cout << "#Dipole size is " << r1-r2 << endl;
-                        //cout << "#Dipole amplitude is " << Amplitude(0, q1, q2) << endl;
-      //                  thrownAway += 1;
-                        continue;
                 }
+                   r = sqrt(pow((X(i)-X(k)),2)+pow(Y(j)-Y(l)),2);
 
-                    // Haw w2 = WilsonLine(k,l)
-                    // Laske w1.MultiplbyHermitianConjugate(w2);
-
-                   sum += Amplitude(0, q1, q2);
-                   //sum += Amplitude(0, quark, q2);
+                   double WF = WFGammaqqbarLongitudinal(double r, double z, double Q, struct* quark_flavor); 
+                   sum += Amplitude(0, quark, q2)*WF*WF;
                 }
             }
         }
     }
     cout << "# total sum = " << sum << endl;
- //   cout << "# i counted " << iWasHere << " values" << endl;
-   // cout << "# threw away " << thrownAway << " values (" << 100.*thrownAway*pow(L,-4) << "%)" << endl;
     return sum;
 }
 
@@ -672,25 +662,34 @@ double IPGlasma::SumAmplitudes()
 // return longitudinal photon polarization wave function
 // source arxiv.org/pdf/hep-ph/0606272
 // quark flavor f: 0=u, 1=d
-double IPGlasma::WFGammaqqbar(double r, int f)
+double IPGlasma::WFGammaqqbarLongitudinal(double r, double z, double Q, struct* quark_flavor)
 { 
     //Assuming r is the transverse size of the dipole
     //constants: e_f (quark charge), e, N_C, 
     //variables: r, Q, z, epsilon
     //epsilon = sqrt(z(1-z)Q^2-m^2_f
 
+    double lattice_spacing = xcoords[1]-xcoords[0];
     // fine structure constant set to 1/137
-    double e = cmath.sqrt(4*cmath.pi*1/137)
-    double m = 0.16e-9 //GeV
+    double e = sqrt(4*M_PI*1/137);
+    double m_f = 0.16e-9; //GeV
     double e_f;
-    if (f){ e_f = -1/3};
-    else { e_f = 2/3};
-
-    double z = 0.5 //GeV^0
-    double QQ = 10 //GeV^2
-    double Delta = 0.5 //GeV
-    double l = a/100 // << a^-1 (a = lattice const)
-    double eps = cmath.sqrt(z*(1-z)*QQ*m_f*m_f)
+    if (f){ e_f = -1/3;}
+    else { e_f = 2/3;}
+
+    double z = 0.5; //GeV^0
+    double Q = sqrt(10); //GeV
+    double Delta = 0.5; //GeV
+    double l = lattice_spacing/100; // << a^-1 (a = lattice const)
+    double eps = sqrt(z*(1-z)*Q*Q*m_f*m_f);
  
-    return e_f*e*cmath.sqrt(NC)*2*Q*z*(1-z)*std::cyl_bessel_k(0, eps*r)0.5/cmath.pi 
+    return (e_f*e*sqrt(NC)*2*Q*Q*z*(1-z)*std::cyl_bessel_k(0, eps*r)*0.5/M_PI);
+}
+
+// TODO: This has to be done inside the loop in SumAmplitudes
+double IPGlasma::DiffractiveCrossSection()
+{
+    
+    SumAmplitudes
+    return 0;
 }
diff --git a/src/ipglasma.hpp b/src/ipglasma.hpp
index 1a979b92fcc70346ef06b5ae3daae602aba72aad..81d30501e5078d03db9748b2af6e0f76e4901515 100644
--- a/src/ipglasma.hpp
+++ b/src/ipglasma.hpp
@@ -61,8 +61,15 @@ public:
     double Amplitude(double xpom, WilsonLine quark, int q2[2]);
     double SumAmplitudes();
     double SumAmplitudes(int l);
-    
+    double WFGammaqqbarLongitudinal(double r, double z, double Q, struct* quark_flavor);
+    struct q_flavor(double e_charge, double mass);
 
+    //TODO: missä paras asettaa arvot?
+    struct q_flavor{
+        double e_charge;
+        double mass;
+    }
+    
 private:
     
     
diff --git a/src/main.cpp b/src/main.cpp
index d130de95f330629653a0c010004cb866302619cb..46f1543b349b7d4a2ed0c96765977a9fa8153225 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -55,66 +55,6 @@ int printDipoleAmplitudes(string fname, IPGlasma ipglasma)
     return 0; 
 }
 
-        int thread_id = omp_get_thread_num();
-        int num_threads = omp_get_num_threads();
-    
-        // Print the number of threads (only from the master thread to avoid too many prints)
-        if (thread_id == 0) {
-            cout << "# Number of threads: " << num_threads << endl;
-        }
-
-        // Print the thread id for debugging (can be limited to a few threads for easier reading)
-        if (thread_id < 5) {  // Change 5 to a larger number if needed
-             cout << "# Thread " << thread_id << " is starting" << endl;
-        }
-
-    } 
-    cout << "# Hello from main.cpp sumAmplitudes!" << endl;
-    double sum = 0.0;
-    int L = 512; //ipglasma->xcoords.size();
-    cout << "# Summing with lattice size L = " << L << endl;
-    //WilsonLine quark;
-    //int thrownAway = 0;
-    //int iWasHere = 0;
-    #pragma omp parallel for reduction(+:sum)//,thrownAway,iWasHere) //collapse(4)
-
-    for(int i =0; i < L; i += 1){ 
-        for(int j =0; j < L; j += 1){ 
-    
-            //#pragma omp parallel for reduction(+:sum) collapse(2)
-            for(int k =0; k < L; k +=1){
-//                cout << "# Thread " << omp_get_thread_num() << " is processing (" << i << ", " << j << ", " << k <<")" << endl;
-                for(int l =0; l < L; l +=1){
-                    int q1[2] = {i,j};
-                    double r1 = ipglasma->DistanceToOrigin(q1);
-                    int q2[2] = {k,l};
-                    //double r1 = ipglasma->DistanceToOrigin(q1);
-                    double r2 = ipglasma->DistanceToOrigin(q2);
-                    if (r1 > 1. && r2 > 1.){
-                       // cout << "#!!Throwing away large dipole separation!" << endl;
-                        //cout << "#Quark coordinate indices " << q1[0] << "," << q1[1] << " and " << q2[0] << "," << q2[1] << endl;
-                        //cout << "#Quark coordinates " << X(q1[0]) << "," << Y(q1[1]) << " and " << X(q2[0]) << "," << Y(q2[1]) << endl;
-                        //cout << "#Dipole size is " << r1-r2 << endl;
-                        //cout << "#Dipole amplitude is " << Amplitude(0, q1, q2) << endl;
-                        //thrownAway += 1;
-                        continue;
-                }
-
-                    // Haw w2 = WilsonLine(k,l)
-                    // Laske w1.MultiplbyHermitianConjugate(w2);
-
-                   sum += ipglasma->Amplitude(0, q1, q2);
-                   //sum += Amplitude(0, quark, q2);
-                   //iWasHere += 1;  
-                }
-            }
-        }
-    }   
-    cout << "# total sum = " << sum << endl;
-    //cout << "# i counted " << iWasHere << " values" << endl;
-    //cout << "# threw away " << thrownAway << " values (" << 100.*thrownAway*pow(L,-4) << "%)" << endl;
-    return sum;
-}
 
 int main(int argc, char* argv[])
 {
@@ -125,20 +65,20 @@ int main(int argc, char* argv[])
     string fname = argv[1];
     IPGlasma ipglasma(fname, 0, BINARY);
 
-    int task_id = 0;
+    if (!argv[2]){ 
+        printDipoleAmplitudes(fname, ipglasma);
+    }
 
+    int task_id = 0;
     if (argc > 3){ 
         task_id= atoi(argv[3]);
         cout << "#Integrating array job #" << task_id << endl;
-        ipglasma->SumAmplitudes(task_id);   
+        ipglasma.SumAmplitudes(task_id);   
     }
 
-    else if (!argv[2]){ 
-        printDipoleAmplitudes(fname, ipglasma);
-    }
     
     else{ 
-        double sum = ipglasma->SumAmplitudes();
+        double sum = ipglasma.SumAmplitudes();
         // Save results
         //std::filesystem::create_directory("output_"+fname);
         ofstream outfile("da_sum_ +" + fname + ".txt");