/**
  \file

  $Id: ems.c,v 1.19 2004/07/14 13:41:20 rauch Exp $

  \srclevel ISO C
  \mtlevel  TBD
  \author   Kevin P. Rauch

  \brief An advanced custom driver adding an explicit lunar sub-integration.

  This custom driver includes the extra kick and energy due to the Sun's 
  tidal attraction on the extended Earth-Moon system,
  using an explicit lunar sub-integration.
  The expansion of the lunar orbit due to tidal dissipation is neglected.
  This file makes use of low-level primitives to advance the lunar
  integration in lock step with the planets, and should be considered an
  "expert" application of the \hnb package. The physical error in
  the lunar model produces relative errors ~1e-6 per orbit; for the
  Earth-Moon berycenter corresponding error (mainly phase) is ~1e-10 per 
  orbit near the initial epoch.

  \return EXIT_SUCCESS on success, else EXIT_FAILURE.

  \see hnb_argv_driver()
*/
#include <math.h>

#include <hnbody/errors.h>
#include <hnbody/expert.h>
#include <hnbody/kepler.h>
#include <hnbody/output.h>


#define IBARY 3
#define IPMIN 1
#define IPMAX 0
#define IMERC 0
#define IMARS 0

/* Main Solar System integration and lunar sub-integration. */
static hnb_data_t *plan, *moon;
static int restarting;

/* Data to manage lunar output. */
static hnb_options_t mopts;
static double moutint, miter0;
static FILE *fmoon=NULL;


void ems_version(FILE *f, const char *prefix)
{
  hnb_util_version(f, "EMS driver", "$Revision: 1.19 $", prefix);
}


static void tidal_kick(double dvbary[], const double xbary[], double mbary,
  double dvsat[], const double xsat[], double msat,
  double dvp[], const double xp[], double mp, double G, double dt)
{
/*
   For an integration containing movement of the barycenter of a
   primary/satellite system, with a dual-integration of the satellite
   itself, computes the kicks on the barycenter and the satellite due to tidal
   perturbations by the body mp of the barycenter integration.
   The two integrations are assumed to proceed in lock step.
   The kicks are stored in dv*[], x*[] is the *bodycentric* position of each
   entity (relative to mp for xbary[]), and the m* are the masses.
*/
  double x0=xbary[0]-xp[0], x1=xbary[1]-xp[1], x2=xbary[2]-xp[2],
	 irb=1/sqrt(x0*x0+x1*x1+x2*x2), dx0, dx1, dx2,
	 alpha, beta, fac, q, gp, gs, mpri=mbary-msat, imb=1/mbary;

  /* Xpri = Xbary + dx. */
  alpha= -msat*imb;
  dx0=alpha*xsat[0]; dx1=alpha*xsat[1]; dx2=alpha*xsat[2];
  q=irb*irb*((x0+0.5*dx0)*dx0+(x1+0.5*dx1)*dx1+(x2+0.5*dx2)*dx2);
  gp=encke_g(q);

  /* Xsat = Xbary + dx. */
  beta=mpri*imb;
  dx0=beta*xsat[0]; dx1=beta*xsat[1]; dx2=beta*xsat[2];
  q=irb*irb*((x0+0.5*dx0)*dx0+(x1+0.5*dx1)*dx1+(x2+0.5*dx2)*dx2);
  gs=encke_g(q);

  /* Effective tidal acceleration on the subsystem barycenter. */
  fac=G*irb*irb*irb*dt;
  alpha=(mp*imb)*fac*(msat*gs+mpri*gp); beta=msat*(mp*imb)*fac*(gs-gp);
  dvbary[0]+=alpha*x0+beta*dx0;
  dvbary[1]+=alpha*x1+beta*dx1;
  dvbary[2]+=alpha*x2+beta*dx2;

  if (dvp) {
    alpha=fac*(msat*gs+mpri*gp); beta=msat*fac*(gs-gp);
    dvp[0]-=alpha*x0+beta*dx0;
    dvp[1]-=alpha*x1+beta*dx1;
    dvp[2]-=alpha*x2+beta*dx2;
  }

  /*
     Tidal acceleration on the satellite. This could be done more
     transparently using a satellite extra_kick(); we minimize overhead
     by inlining it here---assumes Jacobi coordinates for the
     satellite integration! (Multiply by mpri/mbary for bodycentric case.)
  */
  x0+=dx0; x1+=dx1; x2+=dx2; /* Xsat. */
  dx0= -xsat[0]; dx1= -xsat[1]; dx2= -xsat[2];
  irb=1/sqrt(x0*x0+x1*x1+x2*x2);
  q=irb*irb*((x0+0.5*dx0)*dx0+(x1+0.5*dx1)*dx1+(x2+0.5*dx2)*dx2);
  gs=encke_g(q);
  fac= -G*mp*irb*irb*irb*dt;
  alpha=fac*gs; beta=fac*(gs-1);
  dvsat[0]+=alpha*x0+beta*dx0;
  dvsat[1]+=alpha*x1+beta*dx1;
  dvsat[2]+=alpha*x2+beta*dx2;
}

void ems_kick(double (*dv)[3], double t, double (*x)[3],
  const double m[], int nHL, int n, double dt, const hnb_data_t *sys)
{
/*
   Calculate differential Primary+Satellite and {Primary, Satellite}
   acceleration. This routine assumes the barycenter is index IBARY and
   satellite is index 1 in their respective integrations. Optionally correct
   for tidal effects due to the other bodies, and for ephemeris errors in
   Mercury (argperi) and Mars (meananom).
*/
  const double G=sys->G, c=sys->c;
  double *xsat, *dvsat, msat=hnb_masses(moon)[1], xb[3]={0, 0, 0};
  int i, pindex=hnb_extra_index(plan);

  if (pindex) { xsat=moon->x1[1]; dvsat=moon->dv1[1]; }
  else        { xsat=moon->x0[1]; dvsat=moon->dv0[1]; }
  /*
  hnb_shift(moon, 0.5*dt, pindex);
  */

  /* Correct Mercury argperi due to higher order PN corrections. */
  if (IMERC) {
    double *xm=x[IMERC], irm2=1/(xm[0]*xm[0]+xm[1]*xm[1]+xm[2]*xm[2]),
	   *dvm=dv[IMERC], fac=G*m[0]*irm2/c,
	   Gdt=2e-3*(-2*fac*fac+3.5e-8*G*m[0]*irm2*sqrt(irm2))*dt;

    dvm[0]+=Gdt*xm[0]; dvm[1]+=Gdt*xm[1]; dvm[2]+=Gdt*xm[2];
  }

  /* Correct Mars mean motion due to asteroid mean field. */
  if (IMARS) {
    double *xm=x[IMARS], irm=1/sqrt(xm[0]*xm[0]+xm[1]*xm[1]+xm[2]*xm[2]),
	   *dvm=dv[IMARS], Gdt=G*4e-10*m[0]*irm*irm*irm*dt;

    dvm[0]+=Gdt*xm[0]; dvm[1]+=Gdt*xm[1]; dvm[2]+=Gdt*xm[2];
  }

  /* Dominant mass. */
  tidal_kick(dv[IBARY], x[IBARY], m[IBARY], dvsat, xsat, msat,
	     NULL, xb, m[0], G, dt);

  /* Other bodies. */
  for (i=IPMIN+(IPMIN==IBARY); i<=IPMAX; i+=1+(i==IBARY-1))
    tidal_kick(dv[IBARY], x[IBARY], m[IBARY], dvsat, xsat, msat,
	       dv[i], x[i], m[i], G, dt);

  /* Must come after tidal_kick(). */
  hnb_kick(moon, dt, pindex);

  /*
  hnb_shift(moon, 0.5*dt, pindex);
  */
}

void ems_drift(double (*dx)[3], double (*dv)[3], double t, double (*x)[3],
  double (*v)[3], const double m[], int nHL, int n, double dt,
  const hnb_data_t *sys)
{
/*
   This routine exists merely to trigger the drift in the lunar
   sub-integration. It is called twice per planet drift (with half the
   step size), and hence we do double the step size every other call.
*/
  int pindex=hnb_extra_index(plan);
  static int odd=0;

  odd^=1; if (odd) hnb_drift(moon, 2*dt, pindex);
}


void ems_halfstep(hnb_data_t *sys)
{
  /* Copy Earth-Moon vectors, then continue with real halfstep(). */
  void (*step0)(hnb_data_t *)=hnb_user_half_step;

  hnb_setup_output(moon);

  hnb_user_half_step=NULL;
  hnb_halfstep(sys);
  hnb_user_half_step=step0;
}
static void bary_halfstep(hnb_data_t *sys)
{
  /* Convert Earth-Moon vectors to barycentric coordinates. */
  sys->itZ=Barycentric;
  kepler_body2bary(sys->x1, sys->x1, sys->m, sys->n);
  kepler_body2bary(sys->v1, sys->v1, sys->m, sys->n);
}
static void dummy_step(hnb_data_t *sys) { }


static void save_moon(hnb_data_t *ss, double iter)
{
/*
   hnb_save() called on plan; save moon as well (existing file overwritten).
*/
  if (mopts.savfile) {
    hnb_user_save_func=NULL;
    panicn(hnb_save(mopts.savfile, miter0, moon), NULL, "hnb_save",
	   mopts.savfile);
    hnb_user_save_func=save_moon;
    fflush(fmoon);
  }
}
static void rest_moon(hnb_data_t *ss, double key)
{
/*
   hnb_restore() called on plan; restore the moon as well
   (assumes savfile is a valid SaveFile).
*/
  hnb_user_rest_func=NULL;
  panicn(hnb_restore(&miter0, mopts.savfile, moon), NULL, "hnb_restore",
	 mopts.savfile);
  hnb_user_rest_func=rest_moon;

  if (mopts.outfile) {
    if (fmoon) fclose(fmoon);
    hnb_trim_file(mopts.outfile, 1+key/moutint);
    fmoon=fopen(mopts.outfile, "a");
  }
}
static void output_moon(double t, hnb_vector_t xmoon, hnb_vector_t vmoon)
{
/*
   Check if Output interval is due, and write needed data. It is assumed
   that xmoon and vmoon are barycentric (output data line is geocentric).
   !!!! It is assumed that OutputOrder = x1 x2 x3 v1 v2 v3.
*/
  if (mopts.outfile && fmoon) {
    double iter=floor(fabs((t-hnb_tinit(plan))/hnb_stepsize(plan))+0.5);
    int i;

    if (iter>miter0 && (fmod(iter, moutint) == 0.0)) {
      double data[7];

      data[0]=t;
      for (i=0; i<3; i++) {
	data[0+i]=xmoon[1][i]-xmoon[0][i];
	data[3+i]=vmoon[1][i]-vmoon[0][i];
      }
      hnb_write_dataline(data, HNB_OUTPUT, mopts.outdata, 6,
			 mopts.outdig, NULL, fmoon);
    }
    miter0=iter;
  }
}


static double tidal_energy(const double xpri[], double mpri,
  const double xsat[], double msat, const double xbary[],
  const double xp[], double mp, double G)
{
/*
   Robustly compute:  G*mp*mpri*(1/|xbary-xp|-1/|xbary-xp+xpri|)
		     +G*mp*msat*(1/|xbary-xp|-1/|xbary-xp+xsat|)

   even when |xbary-xp| >> |xpri| and |xsat|.
*/
  double x0=xbary[0]-xp[0], x1=xbary[1]-xp[1], x2=xbary[2]-xp[2],
	 dx0, dx1, dx2, irbary=1/sqrt(x0*x0+x1*x1+x2*x2), q, dE=0;

  /* Xpri = Xbary + dx. */
  dx0=xpri[0]; dx1=xpri[1]; dx2=xpri[2];
  q=irbary*irbary*((x0+0.5*dx0)*dx0+(x1+0.5*dx1)*dx1+(x2+0.5*dx2)*dx2);
  dE+=2*G*mp*mpri*irbary*q/(1+2*q+sqrt(1+2*q));

  /* Xsat = Xbary + dx. */
  dx0=xsat[0]; dx1=xsat[1]; dx2=xsat[2];
  q=irbary*irbary*((x0+0.5*dx0)*dx0+(x1+0.5*dx1)*dx1+(x2+0.5*dx2)*dx2);
  dE+=2*G*mp*msat*irbary*q/(1+2*q+sqrt(1+2*q));

  return(dE);
}

static double bary_energy(hnb_vector_t *xps, hnb_vector_t *vps, hnb_data_t *ps)
{
  /* Primary-Satellite total energy; returns *barycentric* xps, vps. */
  void (*step0)(hnb_data_t *);
  double dE, t;

  step0=hnb_user_half_step;
  hnb_user_half_step=bary_halfstep; /* Forces bary xps[] and vps[]. */
  dE=hnb_output(&t, xps, vps, moon);
  hnb_user_half_step=step0;

  /* Check if output is required. */
  output_moon(t, *xps, *vps);

  return(dE);
}

double ems_energy(double t, double (*x)[3], double (*v)[3],
  const double m[], int n, const hnb_data_t *sys)
{
/*
   Compute extra potential energy associated with ems_kick().
   Input vectors assumed to be bodycentric.
*/
  hnb_vector_t xps, vps;
  const double G=sys->G, c=sys->c,
      *xbary=x[IBARY], msat=moon->m[1], xb[3]={0, 0, 0};
  double dE;
  int i;

  dE = bary_energy(&xps, &vps, moon); /* Returns *barycentric* state. */

  /* Mercury higher order PN terms. */
  if (IMERC) {
    double *xm=x[IMERC], irm2=1/(xm[0]*xm[0]+xm[1]*xm[1]+xm[2]*xm[2]),
	   fac=G*m[0]/c;

    dE+=2e-3*m[IMERC]*(-fac*fac*irm2+3.5e-8*G*m[0]*sqrt(irm2));
  }

  /* Mars-asteroid mean field term. */
  if (IMARS) {
    double *xm=x[IMARS];
    dE+=G*4e-10*m[0]*m[IMARS]/sqrt(xm[0]*xm[0]+xm[1]*xm[1]+xm[2]*xm[2]);
  }

  /* Dominant mass. */
  dE+=tidal_energy(xps[0], m[IBARY]-msat, xps[1], msat, xbary, xb, m[0], G);

  /* Other bodies. */
  for (i=IPMIN+(IPMIN==IBARY); i<=IPMAX; i+=1+(i==IBARY-1))
    dE+=tidal_energy(xps[0], m[IBARY]-msat, xps[1], msat, xbary, x[i], m[i], G);

  return(dE);
}

void ems_init(hnb_data_t *ss, const char *initfile)
{
/*
   Store main planet integration location and initialize the lunar
   sub-integration. The lunar integration is pulled along by the planet
   integration via the extra() functions; the use of dummy_step() prevents
   the lunar initialization from taking the initial half-step, which will
   occur automatically with the planets. This function is called immediately
   prior to the initial planet half-step. The initfile argument corresponds to
   the value of the UserInitFile option, if supplied.
*/
  void (*func0)(hnb_data_t *, const char *), (*step0)(hnb_data_t *);
  plan=ss;

  panic(!initfile, "UserInitFile not specified", "ems_init", NULL);
  if (!initfile) exit(EXIT_FAILURE);

  hnb_blank_options(&mopts);
  hnb_read_option_file(&mopts, initfile, stdout);

  hnb_fill_options(&mopts);

  hnb_user_half_step=ems_halfstep;
  hnb_user_save_func=save_moon;
  hnb_user_rest_func=rest_moon;

  /* Save and restore original values. */
  func0=hnb_user_init_func; step0=hnb_user_init_step;

  hnb_user_init_func=NULL;  hnb_user_init_step=dummy_step;
  moon=hnb_opts_init(&mopts, NULL, NULL, NULL, NULL, NULL);
  hnb_user_init_func=func0; hnb_user_init_step=step0;

  /* Set up lunar output, which must be handled within this file. */
  if (mopts.outfile) {
    char *pct=strstr(mopts.outfile, "%d");

    /* Remove %d required by package. */
    if (pct) memmove(pct, pct+2, strlen(pct+1));

    moutint=floor(fabs(mopts.outint/(mopts.outsteps ?
				     1.0 : hnb_stepsize(moon)))+0.5);
    if (!moutint) moutint=1;

    if (!restarting) {
      fmoon=hnb_fopen(mopts.outfile, "w");

      if (mopts.outhead) {
	hnb_header_t h;

        hnb_fill_header(&h, HNB_OUTPUT, mopts.outdata, mopts.outcoord, 0,
	  hnb_G(moon), hnb_c(moon), hnb_M(moon), 0, NULL,
	  hnb_tinit(moon), hnb_stepsize(moon), moutint, 1,
	  mopts.outdig, NULL, mopts.outorder, mopts.notags);

	panic(hnb_write_header(&h, fmoon), NULL, "write_header", mopts.outfile);
      }
      miter0= -1;
    } else
      miter0=1e30;
  }
}


int main(int argc, char *argv[])
{
  hnb_options_t opts;
  hnb_cli_t cli;
  int status;

  hnb_user_init_func=ems_init;

  status=hnb_interface(&opts, &cli, argc, argv, ems_version);
  if (status)
    exit(status == -1 ? EXIT_FAILURE : EXIT_SUCCESS);
  else {
    restarting=(cli.savefile!=NULL);
    status=hnb_opts_driver(&opts, &cli, ems_kick, NULL, ems_drift, NULL,
			   ems_energy);
    hnb_free_options(&opts);
    exit(status);
  }
}