/* TwoPunctures:  File  "TwoPunctures.c"*/

#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <ctype.h>
#include "cctk.h"
#include "cctk_Arguments.h"
#include "cctk_Parameters.h"
#include "TP_utilities.h"
#include "TwoPunctures.h"



/* Swap two variables */
static inline
void swap (CCTK_REAL * restrict const a, CCTK_REAL * restrict const b)
{
  CCTK_REAL const t = *a; *a=*b; *b=t;
}
#undef SWAP
#define SWAP(a,b) (swap(&(a),&(b)))



static
void set_initial_guess(CCTK_POINTER_TO_CONST cctkGH,
                       derivs v)
{
  DECLARE_CCTK_PARAMETERS;

  int nvar = 1, n1 = npoints_A, n2 = npoints_B, n3 = npoints_phi;

  CCTK_REAL *s_x, *s_y, *s_z;
  CCTK_REAL al, A, Am1, be, B, phi, R, r, X;
  CCTK_INT ivar, i, j, k, i3D, indx;
  derivs U;
  FILE *debug_file;

  if (solve_momentum_constraint)
    nvar = 4;

  s_x    =calloc(n1*n2*n3, sizeof(CCTK_REAL));
  s_y    =calloc(n1*n2*n3, sizeof(CCTK_REAL));
  s_z    =calloc(n1*n2*n3, sizeof(CCTK_REAL));
  allocate_derivs (&U, nvar);
  for (ivar = 0; ivar < nvar; ivar++)
    for (i = 0; i < n1; i++)
      for (j = 0; j < n2; j++)
        for (k = 0; k < n3; k++)
        {
          i3D = Index(ivar,i,j,k,1,n1,n2,n3);

          al = Pih * (2 * i + 1) / n1;
          A = -cos (al);
          be = Pih * (2 * j + 1) / n2;
          B = -cos (be);
          phi = 2. * Pi * k / n3;

          /* Calculation of (X,R)*/
          AB_To_XR (nvar, A, B, &X, &R, U);
          /* Calculation of (x,r)*/
          C_To_c (nvar, X, R, &(s_x[i3D]), &r, U);
          /* Calculation of (y,z)*/
          rx3_To_xyz (nvar, s_x[i3D], r, phi, &(s_y[i3D]), &(s_z[i3D]), U);
        }
  Set_Initial_Guess_for_u(cctkGH, n1*n2*n3, v.d0, s_x, s_y, s_z);
  for (ivar = 0; ivar < nvar; ivar++)
    for (i = 0; i < n1; i++)
      for (j = 0; j < n2; j++)
        for (k = 0; k < n3; k++)
        {
          indx = Index(ivar,i,j,k,1,n1,n2,n3);
          v.d0[indx]/=(-cos(Pih * (2 * i + 1) / n1)-1.0);
        }
  Derivatives_AB3 (nvar, n1, n2, n3, v);
  if (do_initial_debug_output && CCTK_MyProc(cctkGH) == 0)
  {
    debug_file=fopen("initial.dat", "w");
    assert(debug_file);
    for (ivar = 0; ivar < nvar; ivar++)
      for (i = 0; i < n1; i++)
        for (j = 0; j < n2; j++)
        {
          al = Pih * (2 * i + 1) / n1;
          A = -cos (al);
          Am1 = A -1.0;
          be = Pih * (2 * j + 1) / n2;
          B = -cos (be);
          phi = 0.0;
          indx = Index(ivar,i,j,0,1,n1,n2,n3);
          U.d0[0] = Am1 * v.d0[indx];        /* U*/
          U.d1[0] = v.d0[indx] + Am1 * v.d1[indx];        /* U_A*/
          U.d2[0] = Am1 * v.d2[indx];        /* U_B*/
          U.d3[0] = Am1 * v.d3[indx];        /* U_3*/
          U.d11[0] = 2 * v.d1[indx] + Am1 * v.d11[indx];        /* U_AA*/
          U.d12[0] = v.d2[indx] + Am1 * v.d12[indx];        /* U_AB*/
          U.d13[0] = v.d3[indx] + Am1 * v.d13[indx];        /* U_AB*/
          U.d22[0] = Am1 * v.d22[indx];        /* U_BB*/
          U.d23[0] = Am1 * v.d23[indx];        /* U_B3*/
          U.d33[0] = Am1 * v.d33[indx];        /* U_33*/
        /* Calculation of (X,R)*/
        AB_To_XR (nvar, A, B, &X, &R, U);
        /* Calculation of (x,r)*/
        C_To_c (nvar, X, R, &(s_x[indx]), &r, U);
        /* Calculation of (y,z)*/
        rx3_To_xyz (nvar, s_x[i3D], r, phi, &(s_y[indx]), &(s_z[indx]), U);
        fprintf(debug_file,
                "%.16g %.16g %.16g %.16g %.16g %.16g %.16g %.16g %.16g %.16g "
                "%.16g %.16g %.16g %.16g %.16g %.16g %.16g %.16g %.16g\n",
                (double)s_x[indx], (double)s_y[indx],
                (double)A,(double)B,
                (double)U.d0[0],
                (double)(-cos(Pih * (2 * i + 1) / n1)-1.0),
                (double)U.d1[0],
                (double)U.d2[0],
                (double)U.d3[0],
                (double)U.d11[0],
                (double)U.d22[0],
                (double)U.d33[0],
                (double)v.d0[indx],
                (double)v.d1[indx],
                (double)v.d2[indx],
                (double)v.d3[indx],
                (double)v.d11[indx],
                (double)v.d22[indx],
                (double)v.d33[indx]
                );
        }
    fprintf(debug_file, "\n\n");
    for (i=n2-10; i<n2; i++)
    {
      CCTK_REAL d;
      indx = Index(0,0,i,0,1,n1,n2,n3);
      d = PunctIntPolAtArbitPosition(0, nvar, n1, n2, n3, v,
              s_x[indx], 0.0, 0.0);
      fprintf(debug_file, "%.16g %.16g\n",
                (double)s_x[indx], (double)d);
    }
    fprintf(debug_file, "\n\n");
    for (i=n2-10; i<n2-1; i++)
    {
      CCTK_REAL d;
      int ip= Index(0,0,i+1,0,1,n1,n2,n3);
      indx = Index(0,0,i,0,1,n1,n2,n3);
      for (j=-10; j<10; j++)
      {
        d = PunctIntPolAtArbitPosition(0, nvar, n1, n2, n3, v,
                s_x[indx]+(s_x[ip]-s_x[indx])*j/10,
                0.0, 0.0);
        fprintf(debug_file, "%.16g %.16g\n",
                (double)(s_x[indx]+(s_x[ip]-s_x[indx])*j/10), (double)d);
      }
    }
    fprintf(debug_file, "\n\n");
    for (i = 0; i < n1; i++)
      for (j = 0; j < n2; j++)
      {
        X = 2*(2.0*i/n1-1.0);
        R = 2*(1.0*j/n2);
        if (X*X+R*R > 1.0)
        {
          C_To_c (nvar, X, R, &(s_x[indx]), &r, U);
          rx3_To_xyz (nvar, s_x[i3D], r, phi, &(s_y[indx]), &(s_z[indx]), U);
          *U.d0  = s_x[indx]*s_x[indx];
          *U.d1  = 2*s_x[indx];
          *U.d2  = 0.0;
          *U.d3 = 0.0;
          *U.d11 = 2.0;
          *U.d22 = 0.0;
          *U.d33 = *U.d12 = *U.d23 = *U.d13 = 0.0;
          C_To_c (nvar, X, R, &(s_x[indx]), &r, U);
          fprintf(debug_file,
                  "%.16g %.16g %.16g %.16g %.16g %.16g %.16g %.16g %.16g %.16g %.16g\n",
                  (double)s_x[indx], (double)r, (double)X, (double)R, (double)U.d0[0],
                  (double)U.d1[0],
                  (double)U.d2[0],
                  (double)U.d3[0],
                  (double)U.d11[0],
                  (double)U.d22[0],
                  (double)U.d33[0]);
        }
      }
    fclose(debug_file);
  }
  free(s_z);
  free(s_y);
  free(s_x);
  free_derivs (&U, nvar);
  /*exit(0);*/
}

/* -------------------------------------------------------------------*/
void
TwoPunctures (CCTK_ARGUMENTS)
{
  DECLARE_CCTK_ARGUMENTS;
  DECLARE_CCTK_PARAMETERS;

  * mp = par_m_plus;
  * mm = par_m_minus;

  enum GRID_SETUP_METHOD { GSM_Taylor_expansion, GSM_evaluation };
  enum GRID_SETUP_METHOD gsm;

  int antisymmetric_lapse, averaged_lapse, pmn_lapse, brownsville_lapse;

  int const nvar = 1, n1 = npoints_A, n2 = npoints_B, n3 = npoints_phi;

  int imin[3], imax[3];
  int const ntotal = n1 * n2 * n3 * nvar;
#if 0
  int percent10 = 0;
#endif
  static CCTK_REAL *F = NULL;
  static derivs u, v;
  CCTK_REAL admMass, M_p, M_m, tmp_p, tmp_m, um, up, new_mass, old_mass;

  if (! F) {
    /* Solve only when called for the first time */
    F = dvector (0, ntotal - 1);
    allocate_derivs (&u, ntotal);
    allocate_derivs (&v, ntotal);

    if (use_sources) {
      CCTK_INFO ("Solving puncture equation for BH-NS/NS-NS system");
    } else {
      CCTK_INFO ("Solving puncture equation for BH-BH system");
    }
    
    /* initialise to 0 */
    for (int j = 0; j < ntotal; j++)
    {
      v.d0[j] = 0.0;
      v.d1[j] = 0.0;
      v.d2[j] = 0.0;
      v.d3[j] = 0.0;
      v.d11[j] = 0.0;
      v.d12[j] = 0.0;
      v.d13[j] = 0.0;
      v.d22[j] = 0.0;
      v.d23[j] = 0.0;
      v.d33[j] = 0.0;
    }
    /* call for external initial guess */
    if (use_external_initial_guess)
    {
      set_initial_guess(cctkGH, v);
    }

    if(!(give_bare_mass)) {
      * mp  = target_M_plus;
      * mm  = target_M_minus;
      
      M_p= target_M_plus;
      M_m= target_M_minus;

      new_mass = 0;
      old_mass = 1.;
      while (new_mass<old_mass) {
	old_mass = (double)*mp + (double)*mm;
	Newton (cctkGH, nvar, n1, n2, n3, v, Newton_tol, 1);

	F_of_v (cctkGH, nvar, n1, n2, n3, v, F, u);

	up  = PunctIntPolAtArbitPosition 
	  (0, nvar, n1, n2, n3, v, par_b, 0., 0.);
	um  = PunctIntPolAtArbitPosition 
	  (0, nvar, n1, n2, n3, v, -par_b, 0., 0.);

	for(int l=0; l<32;l++) {
	  tmp_p=(4*par_b*um+ *mp +4*par_b);
	  tmp_m=(4*par_b*up+ *mm +4*par_b);
	  *mp = *mp - ((*mp*(up+1+M_m/tmp_p)-M_p)/
		       (up+1+(M_m/tmp_p)-(*mp*M_m)/(pow(tmp_p,2))));
	  *mm = *mm - ((*mm*(um+1+M_p/tmp_m)-M_m)/
		       (um+1+(M_p/tmp_m)-(*mm*M_p)/(pow(tmp_m,2))));
	}
	char valbuf_p[100], valbuf_m[100];  
	
	sprintf (valbuf_p,"%f", (float) *mp);
	CCTK_ParameterSet ("par_m_plus", "twopunctures", valbuf_p);
	sprintf (valbuf_m,"%f", (float) *mm);
	CCTK_ParameterSet ("par_m_minus", "twopunctures", valbuf_m);
	new_mass = (double)*mp + (double)*mm;
      }
    } 
    admMass = (*mp + *mm
               - 4*par_b*PunctEvalAtArbitPosition(v.d0, 1, 0, 0, n1, n2, n3));
    CCTK_VInfo (CCTK_THORNSTRING, "ADM mass is %g\n", (double)admMass);

    Newton (cctkGH, nvar, n1, n2, n3, v, Newton_tol, Newton_maxit);
 
    F_of_v (cctkGH, nvar, n1, n2, n3, v, F, u);

    CCTK_VInfo (CCTK_THORNSTRING,
		  "The two puncture masses are %g and %g",
                (double) par_m_minus, (double) par_m_plus);

    /* print out ADM mass, eq.: \Delta M_ADM=2*r*u=4*b*V for A=1,B=0,phi=0 */
    admMass = (*mp + *mm
               - 4*par_b*PunctEvalAtArbitPosition(v.d0, 1, 0, 0, n1, n2, n3));
    CCTK_VInfo (CCTK_THORNSTRING, "ADM mass is %g", (double)admMass);
  }

  if (CCTK_EQUALS(grid_setup_method, "Taylor expansion"))
  {
    gsm = GSM_Taylor_expansion;
  }
  else if (CCTK_EQUALS(grid_setup_method, "evaluation"))
  {
    gsm = GSM_evaluation;
  }
  else
  {
    CCTK_WARN (0, "internal error");
  }

  antisymmetric_lapse = CCTK_EQUALS(initial_lapse, "twopunctures-antisymmetric");
  averaged_lapse = CCTK_EQUALS(initial_lapse, "twopunctures-averaged");
	pmn_lapse = CCTK_EQUALS(initial_lapse, "psi^n");
  if (pmn_lapse)
		CCTK_VInfo(CCTK_THORNSTRING, "Setting initial lapse to psi^%f profile.",
               (double)initial_lapse_psi_exponent);
  brownsville_lapse = CCTK_EQUALS(initial_lapse, "brownsville");
  if (brownsville_lapse)
    CCTK_VInfo(CCTK_THORNSTRING, 
               "Setting initial lapse to a Brownsville-style profile "
               "with exp %f.",
               (double)initial_lapse_psi_exponent);

  CCTK_INFO ("Interpolating result");
  if (CCTK_EQUALS(metric_type, "static conformal")) {
    if (CCTK_EQUALS(conformal_storage, "factor")) {
      *conformal_state = 1;
    } else if (CCTK_EQUALS(conformal_storage, "factor+derivs")) {
      *conformal_state = 2;
    } else if (CCTK_EQUALS(conformal_storage, "factor+derivs+2nd derivs")) {
      *conformal_state = 3;
    }
  } else {
    *conformal_state = 0;
  }

  for (int d = 0; d < 3; ++ d)
  {
    /*
    imin[d] = 0           + (cctk_bbox[2*d  ] ? 0 : cctk_nghostzones[d]);
    imax[d] = cctk_lsh[d] - (cctk_bbox[2*d+1] ? 0 : cctk_nghostzones[d]);
    */
    imin[d] = 0;
    imax[d] = cctk_lsh[d];
  }

#pragma omp parallel for
  for (int k = imin[2]; k < imax[2]; ++k)
  {
    for (int j = imin[1]; j < imax[1]; ++j)
    {
      for (int i = imin[0]; i < imax[0]; ++i)
      {
#if 0
        /* We can't output this when running in parallel */
        if (percent10 != 10*(i+j*cctk_lsh[0]+k*cctk_lsh[0]*cctk_lsh[1]) /
                            (cctk_lsh[0] * cctk_lsh[1] * cctk_lsh[2]))
        {
            percent10 = 10*(i+j*cctk_lsh[0]+k*cctk_lsh[0]*cctk_lsh[1]) /
                           (cctk_lsh[0] * cctk_lsh[1] * cctk_lsh[2]);
            CCTK_VInfo(CCTK_THORNSTRING, "%3d%% done", percent10*10);
        }
#endif

        const int ind = CCTK_GFINDEX3D (cctkGH, i, j, k);

        CCTK_REAL x1, y1, z1;
        x1 = x[ind] - center_offset[0];
        y1 = y[ind] - center_offset[1];
        z1 = z[ind] - center_offset[2];

        /* We implement swapping the x and z coordinates as follows.
           The bulk of the code that performs the actual calculations
           is unchanged.  This code looks only at local variables.
           Before the bulk --i.e., here-- we swap all x and z tensor
           components, and after the code --i.e., at the end of this
           main loop-- we swap everything back.  */
        if (swap_xz) {
          /* Swap the x and z coordinates */
          SWAP (x1, z1);
        }

        CCTK_REAL r_plus
          = sqrt(pow(x1 - par_b, 2) + pow(y1, 2) + pow(z1, 2));
        CCTK_REAL r_minus
          = sqrt(pow(x1 + par_b, 2) + pow(y1, 2) + pow(z1, 2));

        CCTK_REAL U;
        switch (gsm)
        {
        case GSM_Taylor_expansion:
          U = PunctTaylorExpandAtArbitPosition
            (0, nvar, n1, n2, n3, v, x1, y1, z1);
          break;
        case GSM_evaluation:
          U = PunctIntPolAtArbitPosition
            (0, nvar, n1, n2, n3, v, x1, y1, z1);
          break;
        default:
          assert (0);
        }
        r_plus = pow (pow (r_plus, 4) + pow (TP_epsilon, 4), 0.25);
        r_minus = pow (pow (r_minus, 4) + pow (TP_epsilon, 4), 0.25);
        if (r_plus < TP_Tiny)
            r_plus = TP_Tiny;
        if (r_minus < TP_Tiny)
            r_minus = TP_Tiny;
        CCTK_REAL psi1 = 1
          + 0.5 * *mp / r_plus
          + 0.5 * *mm / r_minus + U;
#define EXTEND(M,r) \
          ( M * (3./8 * pow(r, 4) / pow(TP_Extend_Radius, 5) - \
                 5./4 * pow(r, 2) / pow(TP_Extend_Radius, 3) + \
                 15./8 / TP_Extend_Radius))
        if (r_plus < TP_Extend_Radius) {
          psi1 = 1
             + 0.5 * EXTEND(*mp,r_plus)
             + 0.5 * *mm / r_minus + U;
        }
        if (r_minus < TP_Extend_Radius) {
          psi1 = 1
             + 0.5 * EXTEND(*mm,r_minus)
             + 0.5 * *mp / r_plus + U;
        }
        CCTK_REAL static_psi = 1;
        
        CCTK_REAL Aij[3][3];
        BY_Aijofxyz (x1, y1, z1, Aij);

        CCTK_REAL old_alp;
        if (multiply_old_lapse)
            old_alp = alp[ind];

        if ((*conformal_state > 0) || (pmn_lapse) || (brownsville_lapse)) {

          CCTK_REAL xp, yp, zp, rp, ir;
          CCTK_REAL s1, s3, s5;
          CCTK_REAL p, px, py, pz, pxx, pxy, pxz, pyy, pyz, pzz;
          p = 1.0;
          px = py = pz = 0.0;
          pxx = pxy = pxz = 0.0;
          pyy = pyz = pzz = 0.0;

          /* first puncture */
          xp = x1 - par_b;
          yp = y1;
          zp = z1;
          rp = sqrt (xp*xp + yp*yp + zp*zp);
          rp = pow (pow (rp, 4) + pow (TP_epsilon, 4), 0.25);
          if (rp < TP_Tiny)
              rp = TP_Tiny;
          ir = 1.0/rp;

          if (rp < TP_Extend_Radius) {
            ir = EXTEND(1., rp);
          }

          s1 = 0.5* *mp *ir;
          s3 = -s1*ir*ir;
          s5 = -3.0*s3*ir*ir;

          p += s1;

          px += xp*s3;
          py += yp*s3;
          pz += zp*s3;

          pxx += xp*xp*s5 + s3;
          pxy += xp*yp*s5;
          pxz += xp*zp*s5;
          pyy += yp*yp*s5 + s3;
          pyz += yp*zp*s5;
          pzz += zp*zp*s5 + s3;

          /* second puncture */
          xp = x1 + par_b;
          yp = y1;
          zp = z1;
          rp = sqrt (xp*xp + yp*yp + zp*zp);
          rp = pow (pow (rp, 4) + pow (TP_epsilon, 4), 0.25);
          if (rp < TP_Tiny)
              rp = TP_Tiny;
          ir = 1.0/rp;

          if (rp < TP_Extend_Radius) {
            ir = EXTEND(1., rp);
          }

          s1 = 0.5* *mm *ir;
          s3 = -s1*ir*ir;
          s5 = -3.0*s3*ir*ir;

          p += s1;

          px += xp*s3;
          py += yp*s3;
          pz += zp*s3;

          pxx += xp*xp*s5 + s3;
          pxy += xp*yp*s5;
          pxz += xp*zp*s5;
          pyy += yp*yp*s5 + s3;
          pyz += yp*zp*s5;
          pzz += zp*zp*s5 + s3;

          if (*conformal_state >= 1) {
            static_psi = p;
            psi[ind] = static_psi;
          }
          if (*conformal_state >= 2) {
            psix[ind] = px / static_psi;
            psiy[ind] = py / static_psi;
            psiz[ind] = pz / static_psi;
          }
          if (*conformal_state >= 3) {
            psixx[ind] = pxx / static_psi;
            psixy[ind] = pxy / static_psi;
            psixz[ind] = pxz / static_psi;
            psiyy[ind] = pyy / static_psi;
            psiyz[ind] = pyz / static_psi;
            psizz[ind] = pzz / static_psi;
          }

          if (pmn_lapse)
            alp[ind] = pow(p, initial_lapse_psi_exponent);
          if (brownsville_lapse)
            alp[ind] = 2.0/(1.0+pow(p, initial_lapse_psi_exponent));

        } /* if conformal-state > 0 */
          
        puncture_u[ind] = U;

        gxx[ind] = pow (psi1 / static_psi, 4);
        gxy[ind] = 0;
        gxz[ind] = 0;
        gyy[ind] = pow (psi1 / static_psi, 4);
        gyz[ind] = 0;
        gzz[ind] = pow (psi1 / static_psi, 4);

        kxx[ind] = Aij[0][0] / pow(psi1, 2);
        kxy[ind] = Aij[0][1] / pow(psi1, 2);
        kxz[ind] = Aij[0][2] / pow(psi1, 2);
        kyy[ind] = Aij[1][1] / pow(psi1, 2);
        kyz[ind] = Aij[1][2] / pow(psi1, 2);
        kzz[ind] = Aij[2][2] / pow(psi1, 2);

        if (antisymmetric_lapse || averaged_lapse) {
          alp[ind] =
            ((1.0 -0.5* *mp /r_plus -0.5* *mm/r_minus)
            /(1.0 +0.5* *mp /r_plus +0.5* *mm/r_minus));

          if (r_plus < TP_Extend_Radius) {
            alp[ind] =
              ((1.0 -0.5*EXTEND(*mp, r_plus) -0.5* *mm/r_minus)
              /(1.0 +0.5*EXTEND(*mp, r_plus) +0.5* *mm/r_minus));
          }
          if (r_minus < TP_Extend_Radius) {
            alp[ind] =
              ((1.0 -0.5*EXTEND(*mm, r_minus) -0.5* *mp/r_plus)
              /(1.0 +0.5*EXTEND(*mp, r_minus) +0.5* *mp/r_plus));
          }
          
          if (averaged_lapse) {
            alp[ind] = 0.5 * (1.0 + alp[ind]);
          }
        }
        if (multiply_old_lapse)
          alp[ind] *= old_alp;

        if (swap_xz) {
          /* Swap the x and z components of all tensors */
          if (*conformal_state >= 2) {
            SWAP (psix[ind], psiz[ind]);
          }
          if (*conformal_state >= 3) {
            SWAP (psixx[ind], psizz[ind]);
            SWAP (psixy[ind], psiyz[ind]);
          }
          SWAP (gxx[ind], gzz[ind]);
          SWAP (gxy[ind], gyz[ind]);
          SWAP (kxx[ind], kzz[ind]);
          SWAP (kxy[ind], kyz[ind]);
        } /* if swap_xz */

      } /* for i */
    }   /* for j */
  }     /* for k */

  for (int k = imin[2]; k < imax[2]; ++k)
    for (int j = imin[1]; j < imax[1]; ++j)
      for (int i = imin[0]; i < imax[0]; ++i)
        {
          int ijk = CCTK_GFINDEX3D(cctkGH, i, j, k);
          x[ijk] += center_offset[0];
          y[ijk] += center_offset[1];
          z[ijk] += center_offset[2];
        }

  if (use_sources && rescale_sources)
  {
    Rescale_Sources(cctkGH,
                    cctk_lsh[0]*cctk_lsh[1]*cctk_lsh[2],
                    x, y, z,
                    (*conformal_state > 0) ? psi : NULL,
                    gxx, gyy, gzz,
                    gxy, gxz, gyz);
  }

  if (0) {
    /* Keep the result around for the next time */
    free_dvector (F, 0, ntotal - 1);
    free_derivs (&u, ntotal);
    free_derivs (&v, ntotal);
  }
}