//=============================================================================
//
// History: tkcsp (parallel Teukolsky eqn solver implemented by Frans 
//          Pretorius UBC/Penn State Summer School on general relativstic hydro
//
// AMRD/PAMR code for wmap2d
// 
// Lars Andersson, Matthew W Choptuik, 2007--
//
//=============================================================================

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

enum wmap2d_abortvalues {
   ABORT_NONE,
   ABORT_SMAX
};

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

//=============================================================================
// Model parameters
//=============================================================================
real epsilon;

//=============================================================================
// id parameters
//=============================================================================
real amp, cx, cy, sx, sy;
real cr, sr;
real amp1, cr1, sr1, cy1;
real amp2, cr2, sr2, cy2;

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


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

// Pointers to grid functions on specific (active) grid
real *u1_n, *u2_n, *u3_n, *pi1_n, *pi2_n, *pi3_n, *s_n;
real *u1_np1, *u2_np1, *u3_np1, *pi1_np1, *pi2_np1, *pi3_np1, *s_np1;
real *mask;
real *u1_res, *u2_res, *u3_res, *pi1_res, *pi2_res, *pi3_res;
real *temp1, *r, *rh, *phi, *sth, *cth, *cons;

// Vector of global norms
real *g_norms;

// Coordinate, bounding box etc. information
real *x, *y;
// 'size' is total number of points in grid
int  shape[2], ghost_width[4], Nx, Ny, phys_bnd[4], size;
real basebbox[4], bbox[4], dx, dy, dt, t;

// Level associated with grid 
int  g_L;

// Offsets into grid-function arrays for specific grid-functions
int  u1_n_gfn, u2_n_gfn, u3_n_gfn, pi1_n_gfn, pi2_n_gfn, pi3_n_gfn, s_n_gfn;
int  u1_np1_gfn, u2_np1_gfn, u3_np1_gfn, pi1_np1_gfn, pi2_np1_gfn, pi3_np1_gfn, s_np1_gfn;
int  mask_gfn;
int  u1_res_gfn, u2_res_gfn, u3_res_gfn, pi1_res_gfn, pi2_res_gfn, pi3_res_gfn;
int  temp1_gfn, r_gfn, rh_gfn, phi_gfn, sth_gfn, cth_gfn, cons_gfn;

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

    if ((u2_n_gfn = PAMR_get_gfn("u2",PAMR_AMRH,2))<0) AMRD_stop("set_gnfs error",0);
    if ((u2_np1_gfn = PAMR_get_gfn("u2",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);

    if ((u3_n_gfn = PAMR_get_gfn("u3",PAMR_AMRH,2))<0) AMRD_stop("set_gnfs error",0);
    if ((u3_np1_gfn = PAMR_get_gfn("u3",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);

    if ((pi1_n_gfn = PAMR_get_gfn("pi1",PAMR_AMRH,2))<0) AMRD_stop("set_gnfs error",0);
    if ((pi1_np1_gfn = PAMR_get_gfn("pi1",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);

    if ((pi2_n_gfn = PAMR_get_gfn("pi2",PAMR_AMRH,2))<0) AMRD_stop("set_gnfs error",0);
    if ((pi2_np1_gfn = PAMR_get_gfn("pi2",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);

    if ((pi3_n_gfn = PAMR_get_gfn("pi3",PAMR_AMRH,2))<0) AMRD_stop("set_gnfs error",0);
    if ((pi3_np1_gfn = PAMR_get_gfn("pi3",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);

    if ((s_n_gfn = PAMR_get_gfn("s",PAMR_AMRH,2))<0) AMRD_stop("set_gnfs error",0);
    if ((s_np1_gfn = PAMR_get_gfn("s",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);

    if ((u1_res_gfn = PAMR_get_gfn("u1_res",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);
    if ((u2_res_gfn = PAMR_get_gfn("u2_res",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);
    if ((u3_res_gfn = PAMR_get_gfn("u3_res",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);
    if ((pi1_res_gfn = PAMR_get_gfn("pi1_res",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);
    if ((pi2_res_gfn = PAMR_get_gfn("pi2_res",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);
    if ((pi3_res_gfn = PAMR_get_gfn("pi3_res",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);

    if ((r_gfn = PAMR_get_gfn("r",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 ((phi_gfn = PAMR_get_gfn("phi",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);
    if ((sth_gfn = PAMR_get_gfn("sth",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);
    if ((cth_gfn = PAMR_get_gfn("cth",PAMR_AMRH,1))<0) AMRD_stop("set_gnfs error",0);
    if ((cons_gfn = PAMR_get_gfn("cons",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[2],*gfs[PAMR_MAX_GFNS],dx0[2];
   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","");

   u1_n = gfs[u1_n_gfn - 1];
   u1_np1 = gfs[u1_np1_gfn - 1];

   u2_n = gfs[u2_n_gfn - 1];
   u2_np1 = gfs[u2_np1_gfn - 1];

   u3_n = gfs[u3_n_gfn - 1];
   u3_np1 = gfs[u3_np1_gfn - 1];

   pi1_n = gfs[pi1_n_gfn - 1];
   pi1_np1 = gfs[pi1_np1_gfn - 1];

   pi2_n = gfs[pi2_n_gfn - 1];
   pi2_np1 = gfs[pi2_np1_gfn - 1];

   pi3_n = gfs[pi3_n_gfn - 1];
   pi3_np1 = gfs[pi3_np1_gfn - 1];

   s_n = gfs[s_n_gfn - 1];
   s_np1 = gfs[s_np1_gfn - 1];

   pi1_res = gfs[pi1_res_gfn - 1];
   pi2_res = gfs[pi2_res_gfn - 1];
   pi3_res = gfs[pi3_res_gfn - 1];

   mask = gfs[mask_gfn - 1];
   temp1 = gfs[temp1_gfn - 1];
   r = gfs[r_gfn - 1];
   rh = gfs[rh_gfn - 1];
   phi = gfs[phi_gfn - 1];
   sth = gfs[sth_gfn - 1];
   cth = gfs[cth_gfn - 1];
   cons = gfs[cons_gfn - 1];

   x = x0[0]; dx = x[1]-x[0];
   y = x0[1]; dy = y[1]-y[0];

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

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

   Nx = shape[0];
   Ny = shape[1];

   size = Nx*Ny;
}

//=============================================================================
// 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++) for (j = 0; j<Ny; j++) f[i+j*Nx]=0;
   if (phys_bnd[1]==0) for (i = Nx-3; i<Nx; i++) for (j = 0; j<Ny; j++) f[i+j*Nx]=0;
   if (phys_bnd[2]==0) for (i = 0; i<Nx; i++) for (j = 0; j<3; j++) f[i+j*Nx]=0;
   if (phys_bnd[3]==0) for (i = 0; i<Nx; i++) for (j = Ny-3; j<Ny; j++) f[i+j*Nx]=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 indepedent grid functions
//
// Calls RNPL-generated routine: initializer0_
//=============================================================================
void init_tifs()
{
   ldptr();
   initializer0_(r,rh,phi,sth,cth,&Nx,&Ny,x,y,&amp,&cr,&sr);
}

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

//=============================================================================
// Returns 0 to use default mechanism, or is expected to calculate
// the correct initial hierarchy and return 1:
//=============================================================================
int wmap2d_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 wmap2d_var_pre_init(char *pfile)
{
   return;
}

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

   if (my_rank==0)
   {
      printf("===================================================================\n");
      printf("Reading wmap2d parameters:\n\n");
   }
   epsilon = -1.0;

   amp = 0.3;
   cx = 0.0;
   cy = 0.0;
   sx = 0.1;
   sy = 0.1;

   cr = 0.5;
   sr = 0.2;

   amp1 = 0.5;
   cr1 = 0.5;
   sr1 = 0.2;
   cy1 = 0.0;

   amp2 = 0.5;
   cr2 = 0.3;
   sr2 = 0.2;
   cy2 = 1.0;

   consrp = 0;

   AMRD_real_param(pfile,"epsilon",&epsilon,1);

   AMRD_real_param(pfile,"amp",&amp,1);
   AMRD_real_param(pfile,"cx",&cx,1);
   AMRD_real_param(pfile,"cy",&cy,1);
   AMRD_real_param(pfile,"sx",&sx,1);
   AMRD_real_param(pfile,"sy",&sy,1);

   AMRD_real_param(pfile,"cr",&cr,1);
   AMRD_real_param(pfile,"sr",&sr,1);

   AMRD_real_param(pfile,"amp1",&amp1,1);
   AMRD_real_param(pfile,"cr1",&cr1,1);
   AMRD_real_param(pfile,"sr1",&sr1,1);
   AMRD_real_param(pfile,"cy1",&cy1,1);

   AMRD_real_param(pfile,"amp2",&amp1,1);
   AMRD_real_param(pfile,"cr2",&cr1,1);
   AMRD_real_param(pfile,"sr2",&sr1,1);
   AMRD_real_param(pfile,"cy2",&cy1,1);

   AMRD_int_param(pfile,"consrp",&consrp,1);

   AMRD_real_param(pfile,"s_threshold",&s_threshold,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 wmap2d_AMRH_var_clear(void)
{
   ldptr();

   zero(u1_n); zero(u1_np1);
   zero(u2_n); zero(u2_np1);
   zero(u3_n); zero(u3_np1);

   zero(pi1_n); zero(pi1_np1);
   zero(pi2_n); zero(pi2_np1);
   zero(pi3_n); zero(pi3_np1);

   zero(s_n); zero(s_np1);

   init_tifs();

   return;
}

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

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

//=============================================================================
// Initial data for free fields: (at tn = 2)
//
// Calls RNPL-generated routine: initializer1_
//=============================================================================
void wmap2d_free_data(void)
{
   int ltrace = 0;
   ldptr();
   init_tifs();
   if( ltrace ) {
      printf("Nx = %d\n",Nx);
      printf("Ny = %d\n",Ny);
      printf("x[0] = %g\n",x[0]);
      printf("x[Nx-1] = %g\n",x[Nx-1]);
      printf("y[0] = %g\n",y[0]);
      printf("y[Ny-1] = %g\n",y[Ny-1]);
      printf("dx = %g\n",dx);
      printf("dy = %g\n",dy);
      printf("amp = %g\n",amp);
      printf("cr = %g\n",cr);
      printf("sr = %g\n",sr);
   }
   initializer1_(u1_n,u2_n,u3_n,pi1_n,pi2_n,pi3_n,s_n,phi,sth,cth,&Nx,&Ny,&dx,&dy);
}  

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

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

//=============================================================================
// Returns some norm of the residual for the evolution variables,
// called after an evolution iteration:
//
// Calls RNPL-generated routines: res_pi1_, res_pi2_, res_pi3_
//=============================================================================
real wmap2d_evo_residual(void)
{
   real l2norm = 0, l2norm_u1, l2norm_u2, l2norm_u3, 
        l2norm_pi1, l2norm_pi2, l2norm_pi3;
   real CHR_amr_bdy = AMRD_CMASK_OFF;
   int  ltrace = 0, num;

   ldptr();

   if (ltrace && g_L==1) 
   {
        printf("wmap2d_evo_residual\n");
        printf("u1 g_norm=%lf\n",g_norms[u1_n_gfn-1]);
        printf("u2 g_norm=%lf\n",g_norms[u2_n_gfn-1]);
        printf("u3 g_norm=%lf\n",g_norms[u3_n_gfn-1]);
        printf("pi1 g_norm=%lf\n",g_norms[pi1_n_gfn-1]);
        printf("pi2 g_norm=%lf\n",g_norms[pi2_n_gfn-1]);
        printf("pi3 g_norm=%lf\n",g_norms[pi3_n_gfn-1]);
        printf("\n");
   }
   if( ltrace ) printf("Before res_pi1_\n");
   res_pi1_(pi1_res,u1_np1,u1_n,pi1_np1,pi1_n,s_np1,s_n,r,
      rh,&Nx,&Ny,x,y,&dt,&dx,&dy,&epsilon,phys_bnd+1,phys_bnd,phys_bnd+3,
      phys_bnd+2);
   if( ltrace ) printf("After res_pi1_\n");
   res_pi2_(pi2_res,u2_np1,u2_n,pi2_np1,pi2_n,s_np1,s_n,r,
      rh,&Nx,&Ny,x,y,&dt,&dx,&dy,&epsilon,phys_bnd+1,phys_bnd,phys_bnd+3,
      phys_bnd+2);
   if( ltrace ) printf("After res_pi2_\n");
   res_pi3_(pi3_res,u3_np1,u3_n,pi3_np1,pi3_n,s_np1,s_n,r,
      rh,&Nx,&Ny,x,y,&dt,&dx,&dy,&epsilon,phys_bnd+1,phys_bnd,phys_bnd+3,
      phys_bnd+2);
   if( ltrace ) printf("After res_pi3_\n");

   zero_bnd(pi1_res);
   zero_bnd(pi2_res);
   zero_bnd(pi3_res);

   l2norm_pi1 = norm_l2(pi1_res,mask);
   l2norm_pi2 = norm_l2(pi2_res,mask);
   l2norm_pi3 = norm_l2(pi3_res,mask);

   num = 0;
   if (g_norms[pi1_n_gfn-1]>0) { l2norm+=l2norm_pi1/g_norms[pi1_n_gfn-1]; num++; }
   if (g_norms[pi2_n_gfn-1]>0) { l2norm+=l2norm_pi2/g_norms[pi2_n_gfn-1]; num++; }
   if (g_norms[pi3_n_gfn-1]>0) { l2norm+=l2norm_pi3/g_norms[pi3_n_gfn-1]; num++; }

   if (num>0) l2norm/=num;

   if( ltrace && (g_L == 1) ) {
      printf("wmap2d_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 wmap2d_MG_residual(void)
{
   return 0;
}

//=============================================================================
// Performs 1 iteration of the evolution equations 
//
// Calls RNPL generated routine
//=============================================================================
void wmap2d_evolve(int iter)
{
   ldptr();
   update0_(u1_np1,u1_n,u2_np1,u2_n,u3_np1,u3_n,pi1_np1,pi1_n,pi2_np1,pi2_n,
      pi3_np1,pi3_n,s_np1,s_n,r,rh,cons,&Nx,&Ny,x,y,&dt,&dx,&dy,&epsilon,
      phys_bnd+1,phys_bnd,phys_bnd+3,phys_bnd+2);
}

//=============================================================================
// Performs 1 relaxation sweep of the MG equations, and returns an estimate
// of the norm of the residual.
//=============================================================================
real wmap2d_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 wmap2d_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 wmap2d_scale_tre(void)
{
   return;
}

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

   real    s_max;

   valid = PAMR_init_s_iter(L,PAMR_AMRH,0);
   while( valid ) {
      ldptr();
      if( consrp ) {
         cons_rp_(u1_np1,u1_n,u2_np1,u2_n,u3_np1,u3_n,s_np1,s_n,
                        &Nx,&Ny,&t,&consrp,&s_max);
         if( my_rank == 0 ) 
            printf("s_max=%g s_threshold=%g\n",s_max,s_threshold);
      }
      valid = PAMR_next_g();
   }
}
//=============================================================================
// Called before taking a time step on level L.  Implement constraint 
// projection here.
//
// Calls RNPL generated routine: cons_rp_
//=============================================================================
void wmap2d_pre_tstep(int L)
{
   int     ltrace = 0;
   int     valid;

   real    s_max;

   if( ltrace ) {
      printf("wmap2d_pre_step: In routine on level %d, consrp=%d\n",L,consrp);
   }

   valid = PAMR_init_s_iter(L,PAMR_AMRH,0);
   while( valid ) {
      ldptr();
      if( consrp ) {
         cons_rp_(u1_np1,u1_n,u2_np1,u2_n,u3_np1,u3_n,s_np1,s_n,
                        &Nx,&Ny,&t,&consrp,&s_max);
         if( my_rank == 0 ) 
            printf("s_max=%g s_threshold=%g\n",s_max,s_threshold);
         if( s_max > s_threshold ) {
            printf("wmap2d_pre_tstep: Calling MPI_Abort on level %d grid on proc %d\n",L,my_rank);
            printf("wmap2d_pre_tstep: s_max(%g) > s_threshold(%g)\n",
               s_max,s_threshold);
            MPI_Abort(MPI_COMM_WORLD,ABORT_SMAX);
         }
      }
      valid = PAMR_next_g();
   }
}

void wmap2d_post_tstep(int L) {
}

void wmap2d_fill_ex_mask(real *mask, int dim, int *shape, real *bbox, real excised)
{
}

//=============================================================================
int main(int argc, char **argv)
{
   amrd_set_app_pre_tstep_hook(wmap2d_pre_tstep);
   amrd(argc,argv,&wmap2d_id,&wmap2d_var_pre_init,
        &wmap2d_var_post_init, &wmap2d_AMRH_var_clear,
        &wmap2d_free_data, &wmap2d_t0_cnst_data,
        &wmap2d_evo_residual, &wmap2d_MG_residual,
        &wmap2d_evolve, &wmap2d_MG_relax, &wmap2d_L_op, 
        &wmap2d_pre_io_calc, &wmap2d_scale_tre, 
        &wmap2d_post_regrid,&wmap2d_post_tstep,
        &wmap2d_fill_ex_mask,0);
}