#include <assert.h>
#include <stdio.h>

#include "cctk.h"
#include "cctk_Arguments.h"
#include "cctk_Parameters.h"

#include "util_Table.h"



static int const max_num_tracked = 10;



void
PunctureTracker_Init (CCTK_ARGUMENTS)
{
  DECLARE_CCTK_ARGUMENTS;
  DECLARE_CCTK_PARAMETERS;
  
  if (verbose) {
    CCTK_INFO ("Initializing PunctureTracker");
  }
  
  for (int n = 0; n < max_num_tracked; ++ n) {
    if (track[n]) {
      pt_loc_t[n] = cctk_time;
      pt_loc_x[n] = initial_x[n];
      pt_loc_y[n] = initial_y[n];
      pt_loc_z[n] = initial_z[n];
    } else {
      // Initialise to some sensible but unimportant values
      pt_loc_t[n] = 0.0;
      pt_loc_x[n] = 0.0;
      pt_loc_y[n] = 0.0;
      pt_loc_z[n] = 0.0;
    }
    pt_loc_t_p[n] = 0.0;
    pt_loc_x_p[n] = 0.0;
    pt_loc_y_p[n] = 0.0;
    pt_loc_z_p[n] = 0.0;
  }
}



void
PunctureTracker_Track (CCTK_ARGUMENTS)
{
  DECLARE_CCTK_ARGUMENTS;
  DECLARE_CCTK_PARAMETERS;
  
  // Do not track while setting up initial data;
  // time interpolation may fail
  
  if (cctk_iteration == 0) {
    return;
  }
  
  // Some output
  
  if (verbose) {
    CCTK_INFO ("Tracking punctures...");
  }
  
  if (verbose) {
    for (int n = 0; n < max_num_tracked; ++ n) {
      if (track[n]) {
        CCTK_VInfo (CCTK_THORNSTRING,
                    "Puncture #%d is at (%g,%g,%g)",
                    n,
                    (double)pt_loc_x[n],
                    (double)pt_loc_y[n],
                    (double)pt_loc_z[n]);
      }
    }
  }
  
  // Manual time level cycling
  
  for (int n = 0; n < max_num_tracked; ++ n) {
    if (track[n]) {
      pt_loc_t_p[n] = pt_loc_t[n];
      pt_loc_x_p[n] = pt_loc_x[n];
      pt_loc_y_p[n] = pt_loc_y[n];
      pt_loc_z_p[n] = pt_loc_z[n];
      
      pt_loc_t[n] = cctk_time;
    }
  }
  
  // Interpolate
  
  // Dimensions
  int const dim = 3;
  
  // Interpolation operator
  int const operator_handle =
    CCTK_InterpHandle ("Lagrange polynomial interpolation");
  if (operator_handle < 0) {
    CCTK_WARN (CCTK_WARN_ALERT, "Can’t get interpolation handle");
    goto label_return;
  }
  
  // Interpolation parameter table
  int const param_table_handle = Util_TableCreateFromString ("order=4");
  if (param_table_handle < 0) {
    CCTK_WARN (CCTK_WARN_ALERT, "Can't create parameter table");
    goto label_return;
  }
  
  {
    
    // Interpolation coordinate system
    int const coordsys_handle = CCTK_CoordSystemHandle ("cart3d");
    if (coordsys_handle < 0) {
      CCTK_WARN (CCTK_WARN_ALERT, "Can't get coordinate system handle");
      goto label_free_param_table;
    }
    
    // Only processor 0 interpolates
    int const num_points = CCTK_MyProc(cctkGH) == 0 ? max_num_tracked : 0;
    
    // Interpolation coordinates
    assert (dim == 3);
    CCTK_POINTER_TO_CONST interp_coords[3];
    interp_coords[0] = pt_loc_x_p;
    interp_coords[1] = pt_loc_y_p;
    interp_coords[2] = pt_loc_z_p;
    
    // Number of interpolation variables
    int const num_vars = 3;
    
    // Interpolated variables
    assert (num_vars == 3);
    int input_array_indices[3];
    input_array_indices[0] = CCTK_VarIndex ("ADMBase::betax");
    input_array_indices[1] = CCTK_VarIndex ("ADMBase::betay");
    input_array_indices[2] = CCTK_VarIndex ("ADMBase::betaz");
    
    // Interpolation result types
    assert (num_vars == 3);
    CCTK_INT output_array_type_codes[3];
    output_array_type_codes[0] = CCTK_VARIABLE_REAL;
    output_array_type_codes[1] = CCTK_VARIABLE_REAL;
    output_array_type_codes[2] = CCTK_VARIABLE_REAL;
    
    // Interpolation result
    CCTK_REAL pt_betax[max_num_tracked];
    CCTK_REAL pt_betay[max_num_tracked];
    CCTK_REAL pt_betaz[max_num_tracked];
    
    assert (num_vars == 3);
    CCTK_POINTER output_arrays[3];
    output_arrays[0] = pt_betax;
    output_arrays[1] = pt_betay;
    output_arrays[2] = pt_betaz;
    
    // Interpolate
    int const ierr = CCTK_InterpGridArrays
      (cctkGH, dim,
       operator_handle, param_table_handle, coordsys_handle,
       num_points,
       CCTK_VARIABLE_REAL,
       interp_coords,
       num_vars, input_array_indices,
       num_vars, output_array_type_codes, output_arrays);
    if (ierr < 0) {
      CCTK_WARN (CCTK_WARN_ALERT, "Interpolation error");
      goto label_free_param_table;
    }
    
    if (CCTK_MyProc(cctkGH) == 0) {
      
      // Some more output
      
      if (verbose && CCTK_MyProc(cctkGH) == 0) {
        for (int n = 0; n < max_num_tracked; ++ n) {
          if (track[n]) {
          CCTK_VInfo (CCTK_THORNSTRING,
                      "Shift at puncture #%d is at (%g,%g,%g)",
                      n,
                      (double)pt_betax[n],
                      (double)pt_betay[n],
                      (double)pt_betaz[n]);
          }
        }
      }
      
      // Time evolution
      
      for (int n = 0; n < max_num_tracked; ++ n) {
        if (track[n]) {
          CCTK_REAL const dt = pt_loc_t[n] - pt_loc_t_p[n];
          // First order time integrator
          // Michael Koppitz says this works...
          // if it doesn't, we can make it second order accurate
          pt_loc_x[n] = pt_loc_x_p[n] + dt * (- pt_betax[n]);
          pt_loc_y[n] = pt_loc_y_p[n] + dt * (- pt_betay[n]);
          pt_loc_z[n] = pt_loc_z_p[n] + dt * (- pt_betaz[n]);
        }
      }
      
    }
    
    // Broadcast result
    
    CCTK_REAL loc_local[3*max_num_tracked];
    if (CCTK_MyProc(cctkGH) == 0) {
      for (int n = 0; n < max_num_tracked; ++ n) {
        loc_local[                  n] = pt_loc_x[n];
        loc_local[  max_num_tracked+n] = pt_loc_y[n];
        loc_local[2*max_num_tracked+n] = pt_loc_z[n];
      }
    } else {
      for (int n = 0; n < max_num_tracked; ++ n) {
        loc_local[                  n] = 0.0;
        loc_local[  max_num_tracked+n] = 0.0;
        loc_local[2*max_num_tracked+n] = 0.0;
      }
    }
    
    CCTK_REAL loc_global[3*max_num_tracked];
    
    int const handle_sum = CCTK_ReductionHandle ("sum");
    if (handle_sum < 0) {
      CCTK_WARN (CCTK_WARN_ALERT, "Can't get redunction handle");
      goto label_free_param_table;
    }
    
    int const ierr2 = CCTK_ReduceLocArrayToArray1D
      (cctkGH, -1, handle_sum,
       loc_local, loc_global, 3*max_num_tracked, CCTK_VARIABLE_REAL);
    if (ierr2 < 0) {
      CCTK_WARN (CCTK_WARN_ALERT, "Reduction error");
      goto label_free_param_table;
    }
    
    for (int n = 0; n < max_num_tracked; ++ n) {
      pt_loc_x[n] = loc_global[                  n];
      pt_loc_y[n] = loc_global[  max_num_tracked+n];
      pt_loc_z[n] = loc_global[2*max_num_tracked+n];
    }
    
  }
  
  // Done
  
  // Poor man's exception handling
 label_free_param_table:
  Util_TableDestroy (param_table_handle);
  
 label_return:
  return;
}



void
PunctureTracker_SetPositions (CCTK_ARGUMENTS)
{
  DECLARE_CCTK_ARGUMENTS;
  DECLARE_CCTK_PARAMETERS;
  
  for (int n = 0; n < max_num_tracked; ++ n) {
    if (track[n]) {
      
      if (verbose) {
        CCTK_VInfo (CCTK_THORNSTRING,
                    "Setting position of refined region from puncture #%d to (%g,%g,%g)",
                    n,
                    (double)pt_loc_x[n],
                    (double)pt_loc_y[n],
                    (double)pt_loc_z[n]);
      }
      
      // Set position in CarpetRegrid2
      position_x[n] = pt_loc_x[n];
      position_y[n] = pt_loc_y[n];
      position_z[n] = pt_loc_z[n];
      
    }
  }
}