/////////////////////////////////////////////////////////////////////////////
//                                                                         //
//  Calculate and store gravitational moments                              //
//                                                                         //
//  B. Militzer                                     Berkeley 04-17-17      //
//                                                                         //
/////////////////////////////////////////////////////////////////////////////

#include "CMS.h"

double CMS::CalculateMass() const {
  //  Write3(delta.Sum(),lambda.Sum(),zeta.Sum());
  //  Write3(delta.Sum2(),lambda.Sum2(),zeta.Sum2());
  double M = 0.0;
  for(int j=0; j<nl; j++) { // loop over layers

    // note the 2.0 in 2.0/3.0 was probably entered because we use symmetry
    double factor = 2.0/3.0 * delta[j] * cube(lambda[j]) * 2.0 * pi;

    for(int k=0; k<nq/2; k++) {
      M += factor * cube(zeta(j,k)) * wq[k];
    }

  }
  return M; // should be 1 for converged results = total mass divided by planet mass
}

/*
double CMS::CalculateMassOld() const {
  //  Write3(delta.Sum(),lambda.Sum(),zeta.Sum());
  //  Write3(delta.Sum2(),lambda.Sum2(),zeta.Sum2());
  double M = 0.0;
  for(int j=0; j<nl; j++) { // loop over layers
    //    if (delta[j]!=delta[j]) warning("delta[j]",delta[j],j);
    //    if (lambda[j]!=lambda[j]) warning("lambda[j]",lambda[j],j);
    //    double factor = delta[j] * pow(lambda[j],3.0) * 2.0 * pi;
    //    for(int k=0; k<nq; k++) {
    double factor = 2.0 * delta[j] * pow(lambda[j],3.0) * 2.0 * pi;
    for(int k=0; k<nq/2; k++) {
      M += factor * pow(zeta(j,k),3.0) * wq[k];
      //      if (zeta(j,k)!=zeta(j,k)) warning("zeta(j,k)",zeta(j,k),j,k);
    }
  }
  M /= 3.0;
  //  Write(M);
  return M;
}
*/

// Calculate the volume enclosed by the surface of layer 'j'.
double CMS::CalculateSpheroidVolume(const int j) const { 
  double V = 0.0;
  for(int k=0; k<nq/2; k++) {
    V += cube(zeta(j,k)) * wq[k];
  }
  double factor = 2.0/3.0 * cube(lambda[j]) * 2.0 * pi;
  V *= factor;
  //  Write(V);
  return V;
}

void CMS::CalculateMoments() {

  double denom = 0.0;
  
  // calculated the denominator common to J, C and S
  for(int j=0; j<nl; j++) {
    //    double factor = delta[j] * pow(lambda[j],3.0) * 2.0 * pi;
    double factor = 2.0 * delta[j] * pow(lambda[j],3.0) * 2.0 * pi;

    // function: gauss_mult 
    // calculate the volume integral of a function in the axisymmetric
    // case using legendre-gauss.
    //    for(int k=0; k<nq; k++) {
    for(int k=0; k<nq/2.0; k++) {
      denom += factor * pow(zeta(j,k),3.0) * wq[k];
    }
  }

  //  Write(denom); // should be equal 3.0 when converged
  //  Quit("XX");

  // calculate J
  for(int i=0; i<nl; i++) {

    for(int n=1; n<=np; n+=2) { // in axisymmetric case, jTilde=0 for odd n
      jTilde(i,n) = 0.0;
    }
    for(int n=0; n<=np; n+=2) {
      double JPre = -3.0/double(n+3) * delta[i] * pow(lambda[i],3.0)* 2.0*pi; // Sean has 2pi inside gauss_mult()

      double Jnum = 0.0;
      for(int k=0; k<nq/2; k++) { // simplification only works for even n
	//	Write4(i,n,k,pow(zeta(i,k),double(n)+3.0) * Pn_mu(k,n) * wq[k]);
	Jnum += pow(zeta(i,k),double(n)+3.0) * Pn_mu(n,k) * wq[k]; // may want to switch indices of Pn_mu(k,n)
      }
      jTilde(i,n) = 2.0 * JPre * Jnum / denom;
      //      Write5(i,n,JPre,Jnum,jTilde(i,n));
    }
  }

  // calculate J'
  for(int i=0; i<nl; i++) {
    
    for(int n=1; n<=np; n+=2) { // in axisymmetric case, jTildeP=0 for odd n
      jTildeP(i,n) = 0.0;
    }

    for(int n=0; n<=np; n+=2) {
      
      if (n==2) {
	double JPre = -3.0 * delta[i] * pow(lambda[i],3.0)* 2.0*pi; // Sean has 2pi inside gauss_mult()

	double Jnum = 0.0;
	for(int k=0; k<nq; k++) { // this loop needs to run from 0...<nq, not be replaced by 2.0*(0...nq/2)
	  Jnum += log(zeta(i,k)) * Pn_mu(n,k) * wq[k];
	}
	jTildeP(i,n) = JPre * Jnum / denom;

      } else { // n!=2 general case

	double JPre = -3.0/(2.0-n) * delta[i] * pow(lambda[i],3.0)* 2.0*pi;// Sean has 2pi inside gauss_mult()
	double Jnum = 0.0;
	for(int k=0; k<nq/2; k++) { // simplification only works for even n
	  Jnum += pow(zeta(i,k),2.0-n) * Pn_mu(n,k) * wq[k];
	}
	jTildeP(i,n) = 2.0 * JPre * Jnum / denom;
	//	Write5(i,n,JPre,Jnum,jTildeP(i,n));
      }
    }
  }

  // calculate J''
  for(int i=0; i<nl; i++) {
    jPP[i] = 2.0*pi*delta[i] / denom;
  }
  //  Write(jPP);

  CalculateUnnormalizedMoments();
}

void CMS::CalculateUnnormalizedMoments() {
  for(int j=0; j<nl; j++) {

    for(int k=0; k<nq; k++) {
      xi(j,k) = lambda[j]*zeta(j,k); // un-normalized radii
    }

    for(int n=0; n<=np; n++) {
      jReg(j,n)  = pow(lambda[j],n) * jTilde(j,n);
      jRegP(j,n) = jTildeP(j,n) / pow(lambda(j),n+1);
    }

  }

  // calculate total gravitational moments
  for(int n=0; n<=np; n++) {
    jTotal[n]=0.0;
  }

  for(int j=0; j<nl; j++) {
    for(int n=0; n<=np; n++) {
      jTotal[n] += jReg(j,n);
    }
  }

  //  Write(jTotal);
  // print total gravitational moments
  for(int n=0; n<=np; n++) {
    Write2(n,jTotal[n]);
  }

}

void CMS::CalculateNormalizedMoments() {
  for(int j=0; j<nl; j++) {
    
    for(int n=0; n<=np; n++) { // iterate over polynomial order
      jTilde(j,n)  = jReg(j,n)  / pow(lambda[j],n);
      jTildeP(j,n) = jRegP(j,n) / pow(lambda[j],n+1);
    }

  }

  // calculate total gravitational moments
  for(int n=0; n<=np; n++) {
    jTotal[n]=0.0;
  }
  for(int i=0; i<nl; i++) {
    for(int n=0; n<=np; n++) {
      jTotal[n] += jReg(i,n);
    }
  }
}