//=============================================================================
// AMRD/PAMR code for equivarint (rotationally symmetric) wave map
// 
// Lars Andersson, Matthew W Choptuik, 2007--
//
//=============================================================================

#include <stdio.h>
#include <mpi.h>
#include <pamr.h>
#include <amrd.h>
#include <math.h>

enum eqwm_abortvalues {
   ABORT_NONE,
   ABORT_SMAX
};

//=============================================================================
// Extra hook functions
//=============================================================================
void eqwm_pre_tstep(int L);

//=============================================================================
// Model parameters
//=============================================================================
real kappa;

//=============================================================================
// id parameters
//=============================================================================
real amp, r0, rwid, idsignum;

//=============================================================================
// In situ dissipation parameter
//=============================================================================
real epsdis;

//=============================================================================
// Other parameters
//=============================================================================
real s_threshold;
int  consrp;


//=============================================================================
// Convenient, "local" global variables
//=============================================================================

// Pointers to grid functions on specific (active) grid
real *ph_n, *pi_n;
real *ph_np1, *pi_np1;
real *mask;
real *ph_res, *pi_res;
real *rh;

// Vector of global norms
real *g_norms;

// Coordinate, bounding box etc. information
real *r;
// 'size' is total number of points in grid
int  shape[1], ghost_width[2], Nr, phys_bnd[2], size;
real basebbox[2], bbox[2], dr, dt, t;

// Level associated with grid 
int  g_L;

// Offsets into grid-function arrays for specific grid-functions
int  ph_n_gfn, pi_n_gfn;
int  ph_np1_gfn, pi_np1_gfn;
int  mask_gfn;
int  ph_res_gfn, pi_res_gfn;
int  rh_gfn;

//=============================================================================
// call after variables have been defined
//=============================================================================
void set_gfns(void)
{
    if ((ph_n_gfn = PAMR_get_gfn("ph",PAMR_AMRH,2))<0) AMRD_stop("set_gnfs error",0);
    if ((ph_np1_gfn = PAMR_get_gfn("ph",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);

    if ((pi_n_gfn = PAMR_get_gfn("pi",PAMR_AMRH,2))<0) AMRD_stop("set_gnfs error",0);
    if ((pi_np1_gfn = PAMR_get_gfn("pi",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);

    if ((ph_res_gfn = PAMR_get_gfn("ph_res",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);
    if ((pi_res_gfn = PAMR_get_gfn("pi_res",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);

    if ((rh_gfn = PAMR_get_gfn("rh",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);

    if ((mask_gfn = PAMR_get_gfn("cmask",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);

    g_norms = AMRD_get_global_norms();
}

//=============================================================================
// call with valid iter to set up globals:
//=============================================================================
void ldptr(void)
{
   int rank,dim,ngfs;
   real *x0[1],*gfs[PAMR_MAX_GFNS],dr0[1];
   static int first = 1;


   if (first) 
   {
      first = 0; 
      set_gfns();
      PAMR_get_global_bbox(basebbox);
   }

   if (!(PAMR_get_g_attribs(&rank,&dim,shape,bbox,ghost_width,&t,&ngfs,x0,gfs))) 
      AMRD_stop("ldptr: PAMR_get_g_attribs failed\n","");

   ph_n = gfs[ph_n_gfn - 1];
   ph_np1 = gfs[ph_np1_gfn - 1];

   pi_n = gfs[pi_n_gfn - 1];
   pi_np1 = gfs[pi_np1_gfn - 1];

   ph_res = gfs[ph_res_gfn - 1];
   pi_res = gfs[pi_res_gfn - 1];

   mask = gfs[mask_gfn - 1];
   rh = gfs[rh_gfn - 1];

   r = x0[0]; dr = r[1]-r[0];

   PAMR_get_g_level(&g_L);
   PAMR_get_g_t(&t);
   PAMR_get_dxdt(g_L,dr0,&dt);

   if ((bbox[0]-basebbox[0])<dr/2) phys_bnd[0]=1; else phys_bnd[0]=0;
   if ((basebbox[1]-bbox[1])<dr/2) phys_bnd[1]=1; else phys_bnd[1]=0;

   Nr = shape[0];

   size = Nr;
}

//=============================================================================
// utility routines
//=============================================================================
void const_f(real *f, real c)
{
   int i;

   for (i = 0; i<size; i++) f[i]=c;
}

void zero(real *f)
{
   const_f(f,0);
}

void zero_bnd(real *f)
{
   int i,j;

   if (phys_bnd[0]==0) for (i = 0; i<3; i++) f[i]=0;
   if (phys_bnd[1]==0) for (i = Nr-3; i<Nr; i++) f[i]=0;

}

real norm_l2(real *f, real* chr)
{
   int i;
   real norm = 0;
   int sum = 0;

   for (i = 0; i<size; i++) 
      if (chr[i]==AMRD_CMASK_ON) { sum++; norm+=f[i]*f[i]; }

   if (!sum) sum = 1;
   return (sqrt(norm/sum));
}

//=============================================================================
// Initialize time independent grid functions
//
// Calls RNPL-generated routine: initializer0_
//=============================================================================
void init_tifs()
{
   ldptr();
   initializer0_(rh,&Nr,r);
}

//=============================================================================
// Routines required by amrd:
//=============================================================================

//=============================================================================
// Returns 0 to use default mechanism, or is expected to calculate
// the correct initial hierarchy and return 1:
//=============================================================================
int eqwm_id(void)
{
   return 0;
}

//=============================================================================
// Sets custom parameters, variables, etc. Split up into two segments,
// one called before the pamr context is initialized and standard
// parameters are read, and the other afterwards
//=============================================================================
void eqwm_var_pre_init(char *pfile)
{
   return;
}

#define DEBUG_MPI 1
void eqwm_var_post_init(char *pfile)
{

   if (my_rank==0)
   {
      printf("===================================================================\n");
      printf("Reading eqwm parameters:\n\n");
   }
   AMRD_real_param(pfile,"kappa",&kappa,1);

   AMRD_real_param(pfile,"amp",&amp,1);
   AMRD_real_param(pfile,"r0",&r0,1);
   AMRD_real_param(pfile,"rwid",&rwid,1);
   AMRD_real_param(pfile,"idsignum",&idsignum,1);
   AMRD_real_param(pfile,"epsdis",&epsdis,1);

   if (DEBUG_MPI) MPI_Errhandler_set(MPI_COMM_WORLD,MPI_ERRORS_RETURN);

   if (my_rank==0) printf("===================================================================\n");
   return;
}


//=============================================================================
// Sets all variables to their 'zero' values:
//=============================================================================
void eqwm_AMRH_var_clear(void)
{
   ldptr();

   zero(ph_n); zero(ph_np1);
   zero(pi_n); zero(pi_np1);

   init_tifs();

   return;
}

//=============================================================================
// Called after regridding (to re-initialize time-independent grid functions)
//=============================================================================

void eqwm_post_regrid(void)
{
   ldptr();
   init_tifs();
}

//=============================================================================
// Initial data for free fields: (at tn = 2)
//
// Calls RNPL-generated routine: ,initializer1_
//=============================================================================
void eqwm_free_data(void)
{
   int ltrace = 0;
   ldptr();
   init_tifs();
   if( ltrace ) {
      printf("Nr = %d\n",Nr);
      printf("r[0] = %g\n",r[0]);
      printf("r[Nr-1] = %g\n",r[Nr-1]);
      printf("dr = %g\n",dr);
      printf("amp = %g\n",amp);
      printf("rwid = %g\n",rwid);
      printf("r0 = %g\n",r0);
      printf("idsignum = %g\n",idsignum);
   }
   initializer1_(ph_n,pi_n,rh,&Nr,r,&dr,&amp,&idsignum,&r0,&rwid,phys_bnd);
}  

//=============================================================================
// Initial constraint data --- called after each MG iteration:
//=============================================================================
void eqwm_t0_cnst_data(void)
{
}

//=============================================================================
// Calculations prior to saving info to disk.
//=============================================================================
void eqwm_pre_io_calc(void)
{
}

//=============================================================================
// Returns some norm of the residual for the evolution variables,
// called after an evolution iteration:
//
// Calls RNPL-generated routines: res_ph_, res_pi_
//=============================================================================
real eqwm_evo_residual(void)
{
   real l2norm = 0, l2norm_ph, l2norm_pi;
   real CHR_amr_bdy = AMRD_CMASK_OFF;
   int  ltrace = 0, num;

   ldptr();

   if (ltrace && g_L==1) 
   {
        printf("eqwm_evo_residual\n");
        printf("ph g_norm=%lf\n",g_norms[ph_n_gfn-1]);
        printf("pi g_norm=%lf\n",g_norms[pi_n_gfn-1]);
        printf("\n");
   }
   res_ph_(ph_res,ph_np1,ph_n,pi_np1,pi_n,rh,&Nr,&dr,&dt,&epsdis,
      phys_bnd,phys_bnd+1);
   res_pi_(pi_res,ph_np1,ph_n,pi_np1,pi_n,rh,&Nr,r,&dr,&dt,&epsdis,&kappa,
      phys_bnd,phys_bnd+1);

   zero_bnd(ph_res);
   zero_bnd(pi_res);

   l2norm_ph = norm_l2(ph_res,mask);
   l2norm_pi = norm_l2(pi_res,mask);

   num = 0;
   if (g_norms[ph_n_gfn-1]>0) { l2norm+=l2norm_ph/g_norms[ph_n_gfn-1]; num++; }
   if (g_norms[pi_n_gfn-1]>0) { l2norm+=l2norm_pi/g_norms[pi_n_gfn-1]; num++; }

   if (num>0) l2norm/=num;

   if( ltrace && (g_L == 1) ) {
      printf("eqwm_evo_residual: Returning l2norm=%g\n",l2norm);
   }
   return l2norm;
}

//=============================================================================
// Returns some norm of the residual for the MG variables, *AND* 
// stores the point-wise residual in "f_res" for each MG variable "f" (for
// computing new RHS's)
//=============================================================================
real eqwm_MG_residual(void)
{
   return 0;
}

//=============================================================================
// Performs 1 iteration of the evolution equations 
//
// Calls RNPL generated routine
//=============================================================================
void eqwm_evolve(int iter)
{
   ldptr();
   update0_(ph_np1,ph_n,pi_np1,pi_n,rh,&Nr,r,&dr,&dt,&epsdis,&kappa,
      phys_bnd,phys_bnd+1); 
}

//=============================================================================
// Performs 1 relaxation sweep of the MG equations, and returns an estimate
// of the norm of the residual.
//=============================================================================
real eqwm_MG_relax(void)
{
   return 0;
}

//=============================================================================
// Computes the differential operator L acting on the current grid,
// in a region specified by the grid function "mask". Stores the result
// in "f_lop" for each MG variable "f"
//=============================================================================
void eqwm_L_op(void)
{
   return;
}

//=============================================================================
// Called after calculating the TRE for all variables
//
// Here, we scale TRE by [csx*csy*csz]^n
//
// NOTE: USING AMRD GLOBAL POINTER VARIABLES HERE (so that we don't need
// to keep track of which vars are used for the TRE)
//=============================================================================
void eqwm_scale_tre(void)
{
   return;
}

//=============================================================================
// Called before taking a time step on level L.  Implement constraint 
// projection here.
//
// Calls RNPL generated routine: cons_rp_
//=============================================================================
void eqwm_pre_tstep(int L)
{
   return;
}

void eqwm_post_tstep(int L) 
{
   return;
}

void eqwm_fill_ex_mask(real *mask, int dim, int *shape, real *bbox, real excised)
{
   return;
}

//=============================================================================
int main(int argc, char **argv)
{
   amrd_set_app_pre_tstep_hook(eqwm_pre_tstep);
   amrd(argc,argv,&eqwm_id,&eqwm_var_pre_init,
        &eqwm_var_post_init, &eqwm_AMRH_var_clear,
        &eqwm_free_data, &eqwm_t0_cnst_data,
        &eqwm_evo_residual, &eqwm_MG_residual,
        &eqwm_evolve, &eqwm_MG_relax, &eqwm_L_op, 
        &eqwm_pre_io_calc, &eqwm_scale_tre, 
        &eqwm_post_regrid,&eqwm_post_tstep,
        &eqwm_fill_ex_mask,0);
}