// Python - C extension module for shallow_water.py
//
// To compile (Python2.3):
//  gcc -c domain_ext.c -I/usr/include/python2.3 -o domain_ext.o -Wall -O 
//  gcc -shared domain_ext.o  -o domain_ext.so  
//
// or use python compile.py
//
// See the module shallow_water.py
//
//
// Ole Nielsen, GA 2004
    
    
#include "Python.h"
#include "Numeric/arrayobject.h"
#include "math.h"

//Shared code snippets
#include "util_ext.h"


// Computational function for rotation
int _rotate(double *q, double n1, double n2) {
  /*Rotate the momentum component q (q[1], q[2])
    from x,y coordinates to coordinates based on normal vector (n1, n2).
    
    Result is returned in array 3x1 r     
    To rotate in opposite direction, call rotate with (q, n1, -n2)
    
    Contents of q are changed by this function */    


  double q1, q2;
  
  //Shorthands
  q1 = q[1];  //uh momentum
  q2 = q[2];  //vh momentum

  //Rotate
  q[1] =  n1*q1 + n2*q2;
  q[2] = -n2*q1 + n1*q2; 

  return 0;
}  


	
// Computational function for flux computation (using stage w=z+h)
int flux_function(double *q_left, double *q_right, 
		  double z_left, double z_right, 
		  double n1, double n2, 
		  double epsilon, double g, 
		  double *edgeflux, double *max_speed) {
  
  /*Compute fluxes between volumes for the shallow water wave equation
    cast in terms of the 'stage', w = h+z using 
    the 'central scheme' as described in
    
    Kurganov, Noelle, Petrova. 'Semidiscrete Central-Upwind Schemes For
    Hyperbolic Conservation Laws and Hamilton-Jacobi Equations'.
    Siam J. Sci. Comput. Vol. 23, No. 3, pp. 707-740.
    
    The implemented formula is given in equation (3.15) on page 714 
  */
       
  int i;
       
  double w_left, h_left, uh_left, vh_left, u_left;
  double w_right, h_right, uh_right, vh_right, u_right;
  double s_min, s_max, soundspeed_left, soundspeed_right;
  double denom, z;
  double q_left_copy[3], q_right_copy[3];    
  double flux_right[3], flux_left[3];
  
  //Copy conserved quantities to protect from modification
  for (i=0; i<3; i++) {
    q_left_copy[i] = q_left[i];
    q_right_copy[i] = q_right[i];
  } 
    
  //Align x- and y-momentum with x-axis
  _rotate(q_left_copy, n1, n2);
  _rotate(q_right_copy, n1, n2);    

  z = (z_left+z_right)/2; //Take average of field values

  //Compute speeds in x-direction
  w_left = q_left_copy[0];              // h+z
  h_left = w_left-z;
  uh_left = q_left_copy[1];

  if (h_left < epsilon) {
    h_left = 0.0;  //Could have been negative
    u_left = 0.0;
  } else { 
    u_left = uh_left/h_left;
  }
  
  w_right = q_right_copy[0];
  h_right = w_right-z;
  uh_right = q_right_copy[1];

  if (h_right < epsilon) {
    h_right = 0.0; //Could have been negative
    u_right = 0.0;
  } else { 
    u_right = uh_right/h_right;
  }

  //Momentum in y-direction              
  vh_left  = q_left_copy[2];
  vh_right = q_right_copy[2];   
       

  //Maximal and minimal wave speeds
  soundspeed_left  = sqrt(g*h_left); 
  soundspeed_right = sqrt(g*h_right);
    
  s_max = max(u_left+soundspeed_left, u_right+soundspeed_right);
  if (s_max < 0.0) s_max = 0.0; 
  
  s_min = min(u_left-soundspeed_left, u_right-soundspeed_right);
  if (s_min > 0.0) s_min = 0.0;   
  
  //Flux formulas  
  flux_left[0] = u_left*h_left;
  flux_left[1] = u_left*uh_left + 0.5*g*h_left*h_left;
  flux_left[2] = u_left*vh_left;
  
  flux_right[0] = u_right*h_right;
  flux_right[1] = u_right*uh_right + 0.5*g*h_right*h_right;
  flux_right[2] = u_right*vh_right;
    

  //Flux computation    
  denom = s_max-s_min;
  if (denom == 0.0) {
    for (i=0; i<3; i++) edgeflux[i] = 0.0;
    *max_speed = 0.0;
  } else {    
    for (i=0; i<3; i++) {
      edgeflux[i] = s_max*flux_left[i] - s_min*flux_right[i];
      edgeflux[i] += s_max*s_min*(q_right_copy[i]-q_left_copy[i]);
      edgeflux[i] /= denom;
    }  
        
    //Maximal wavespeed
    *max_speed = max(fabs(s_max), fabs(s_min));
    
    //Rotate back       
    _rotate(edgeflux, n1, -n2);
  }
  return 0;
}
	
void _manning_friction(double g, double eps, int N, 
		       double* w, double* uh, double* vh, 
		       double* eta, double* xmom, double* ymom) {     

  int k;
  double S;
  
  for (k=0; k<N; k++) {
    if (w[k] >= eps) {
      S = -g * eta[k]*eta[k] * sqrt((uh[k]*uh[k] + vh[k]*vh[k]));
      S /= pow(w[k], 7.0/3);      

      //Update momentum
      xmom[k] += S*uh[k];
      ymom[k] += S*vh[k];
    }
  }
  
}
            

	
///////////////////////////////////////////////////////////////////
// Gateways to Python

PyObject *gravity(PyObject *self, PyObject *args) {
  //
  //  gravity(g, h, v, x, xmom, ymom)
  //
  
  
  PyArrayObject *h, *v, *x, *xmom, *ymom;
  int k, i, N, k3, k6;
  double g, avg_h, zx, zy;
  double x0, y0, x1, y1, x2, y2, z0, z1, z2;
    
  if (!PyArg_ParseTuple(args, "dOOOOO",
			&g, &h, &v, &x, 
			&xmom, &ymom))  
    return NULL;  

  N = h -> dimensions[0];
  for (k=0; k<N; k++) {
    k3 = 3*k;  // base index
    k6 = 6*k;  // base index    
    
    avg_h = 0.0;
    for (i=0; i<3; i++) {
      avg_h += ((double *) h -> data)[k3+i];
    }	
    avg_h /= 3;
        
    
    //Compute bed slope
    x0 = ((double*) x -> data)[k6 + 0];
    y0 = ((double*) x -> data)[k6 + 1];    
    x1 = ((double*) x -> data)[k6 + 2];
    y1 = ((double*) x -> data)[k6 + 3];        
    x2 = ((double*) x -> data)[k6 + 4];
    y2 = ((double*) x -> data)[k6 + 5];           


    z0 = ((double*) v -> data)[k3 + 0];
    z1 = ((double*) v -> data)[k3 + 1];
    z2 = ((double*) v -> data)[k3 + 2];        

    _gradient(x0, y0, x1, y1, x2, y2, z0, z1, z2, &zx, &zy);

    //Update momentum
    ((double*) xmom -> data)[k] += -g*zx*avg_h;
    ((double*) ymom -> data)[k] += -g*zy*avg_h;        
  }
    
  return Py_BuildValue("");
}


PyObject *manning_friction(PyObject *self, PyObject *args) {
  //
  // manning_friction(g, eps, w, uh, vh, eta, xmom_update, ymom_update)
  //
  
  
  PyArrayObject *w, *uh, *vh, *eta, *xmom, *ymom;
  int N;
  double g, eps;
    
  if (!PyArg_ParseTuple(args, "ddOOOOOO",
			&g, &eps, &w, &uh, &vh, &eta, 
			&xmom, &ymom))  
    return NULL;  

  N = w -> dimensions[0];    
  _manning_friction(g, eps, N,
		    (double*) w -> data,
		    (double*) uh -> data, 
		    (double*) vh -> data, 
		    (double*) eta -> data,
		    (double*) xmom -> data, 
		    (double*) ymom -> data);

  return Py_BuildValue("");
}		    

PyObject *rotate(PyObject *self, PyObject *args, PyObject *kwargs) {
  //
  // r = rotate(q, normal, direction=1)
  //
  // Where q is assumed to be a Float numeric array of length 3 and
  // normal a Float numeric array of length 2.

  
  PyObject *Q, *Normal;
  PyArrayObject *q, *r, *normal;
  
  static char *argnames[] = {"q", "normal", "direction", NULL};
  int dimensions[1], i, direction=1;
  double n1, n2;

  // Convert Python arguments to C  
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO|i", argnames, 
				   &Q, &Normal, &direction))  
    return NULL;  

  //Input checks (convert sequences into numeric arrays)
  q = (PyArrayObject *) 
    PyArray_ContiguousFromObject(Q, PyArray_DOUBLE, 0, 0);
  normal = (PyArrayObject *) 
    PyArray_ContiguousFromObject(Normal, PyArray_DOUBLE, 0, 0);
  
  //Allocate space for return vector r (don't DECREF) 
  dimensions[0] = 3;
  r = (PyArrayObject *) PyArray_FromDims(1, dimensions, PyArray_DOUBLE);

  //Copy
  for (i=0; i<3; i++) {
    ((double *) (r -> data))[i] = ((double *) (q -> data))[i];  
  }
  
  //Get normal and direction
  n1 = ((double *) normal -> data)[0]; 
  n2 = ((double *) normal -> data)[1];   
  if (direction == -1) n2 = -n2;

  //Rotate
  _rotate((double *) r -> data, n1, n2);

  //Release numeric arrays
  Py_DECREF(q);    
  Py_DECREF(normal);
        
  //return result using PyArray to avoid memory leak
  return PyArray_Return(r);
}    





PyObject *compute_fluxes(PyObject *self, PyObject *args) {
  /*Compute all fluxes and the timestep suitable for all volumes
    in domain.

    Compute total flux for each conserved quantity using "flux_function"

    Fluxes across each edge are scaled by edgelengths and summed up
    Resulting flux is then scaled by area and stored in
    explicit_update for each of the three conserved quantities 
    level, xmomentum and ymomentum

    The maximal allowable speed computed by the flux_function for each volume
    is converted to a timestep that must not be exceeded. The minimum of
    those is computed as the next overall timestep.
    
    Python call:
    domain.timestep = compute_fluxes(timestep, 
                                     domain.epsilon, 
				     domain.g,
                                     domain.neighbours,
                                     domain.neighbour_edges,
                                     domain.normals,
                                     domain.edgelengths,                       
                                     domain.radii,
                                     domain.areas,
                                     Level.edge_values, 
                                     Xmom.edge_values, 
                                     Ymom.edge_values,  
                                     Bed.edge_values,   
                                     Level.boundary_values,
                                     Xmom.boundary_values,
                                     Ymom.boundary_values,
                                     Level.explicit_update,
                                     Xmom.explicit_update,
                                     Ymom.explicit_update)
        

    Post conditions:
      domain.explicit_update is reset to computed flux values
      domain.timestep is set to the largest step satisfying all volumes. 

            
  */

  
  PyArrayObject *neighbours, *neighbour_edges,
    *normals, *edgelengths, *radii, *areas,
    *level_edge_values, 
    *xmom_edge_values, 
    *ymom_edge_values,  
    *bed_edge_values,   
    *level_boundary_values,
    *xmom_boundary_values,
    *ymom_boundary_values,
    *level_explicit_update,
    *xmom_explicit_update,
    *ymom_explicit_update;

    
  //Local variables  
  double timestep, max_speed, epsilon, g;
  double normal[2], ql[3], qr[3], zl, zr;
  double flux[3], edgeflux[3]; //Work arrays for summing up fluxes

  int number_of_elements, k, i, j, m, n;
  int ki, nm, ki2; //Index shorthands
  
 
  // Convert Python arguments to C  
  if (!PyArg_ParseTuple(args, "dddOOOOOOOOOOOOOOOO",
			&timestep,
			&epsilon,
			&g,
			&neighbours, 
			&neighbour_edges,
			&normals, 
			&edgelengths, &radii, &areas,
			&level_edge_values, 
			&xmom_edge_values, 
			&ymom_edge_values,  
			&bed_edge_values,   
			&level_boundary_values,
			&xmom_boundary_values,
			&ymom_boundary_values,
			&level_explicit_update,
			&xmom_explicit_update,
			&ymom_explicit_update)) {
    PyErr_SetString(PyExc_RuntimeError, "Input arguments failed");
    return NULL;
  }

  number_of_elements = level_edge_values -> dimensions[0];
  
    
  for (k=0; k<number_of_elements; k++) {

    //Reset work array
    for (j=0; j<3; j++) flux[j] = 0.0;
			  
    //Loop through neighbours and compute edge flux for each
    for (i=0; i<3; i++) {
      ki = k*3+i;
      ql[0] = ((double *) level_edge_values -> data)[ki];
      ql[1] = ((double *) xmom_edge_values -> data)[ki];
      ql[2] = ((double *) ymom_edge_values -> data)[ki];            
      zl =    ((double *) bed_edge_values -> data)[ki];                  
      
      //Quantities at neighbour on nearest face
      n = ((int *) neighbours -> data)[ki];
      if (n < 0) {
	m = -n-1; //Convert negative flag to index
	qr[0] = ((double *) level_boundary_values -> data)[m];	
	qr[1] = ((double *) xmom_boundary_values -> data)[m];	
	qr[2] = ((double *) ymom_boundary_values -> data)[m];	
	zr = zl; //Extend bed elevation to boundary
      } else {    
	m = ((int *) neighbour_edges -> data)[ki];
	
	nm = n*3+m;	
	qr[0] = ((double *) level_edge_values -> data)[nm];
	qr[1] = ((double *) xmom_edge_values -> data)[nm];
	qr[2] = ((double *) ymom_edge_values -> data)[nm];            
	zr =    ((double *) bed_edge_values -> data)[nm];                  
      }
      
      // Outward pointing normal vector   
      // normal = domain.normals[k, 2*i:2*i+2]
      ki2 = 2*ki; //k*6 + i*2
      normal[0] = ((double *) normals -> data)[ki2];
      normal[1] = ((double *) normals -> data)[ki2+1];      

      //Edge flux computation
      flux_function(ql, qr, zl, zr, 
		    normal[0], normal[1],
		    epsilon, g, 
		    edgeflux, &max_speed);

		    
      //flux -= edgeflux * edgelengths[k,i]
      for (j=0; j<3; j++) { 
	flux[j] -= edgeflux[j]*((double *) edgelengths -> data)[ki];
      }
      
      //Update timestep
      //timestep = min(timestep, domain.radii[k]/max_speed)
      if (max_speed > epsilon) {
	timestep = min(timestep, ((double *) radii -> data)[k]/max_speed);
      }   
    } // end for i
    
    //Normalise by area and store for when all conserved
    //quantities get updated
    // flux /= areas[k]
    for (j=0; j<3; j++) { 
      flux[j] /= ((double *) areas -> data)[k];
    }

    ((double *) level_explicit_update -> data)[k] = flux[0];
    ((double *) xmom_explicit_update -> data)[k] = flux[1];
    ((double *) ymom_explicit_update -> data)[k] = flux[2];        

  } //end for k

  return Py_BuildValue("d", timestep);
}    








//////////////////////////////////////////	
// Method table for python module
static struct PyMethodDef MethodTable[] = {
  /* The cast of the function is necessary since PyCFunction values
   * only take two PyObject* parameters, and rotate() takes
   * three.
   */
  
  {"rotate", (PyCFunction)rotate, METH_VARARGS | METH_KEYWORDS, "Print out"},
  {"compute_fluxes", compute_fluxes, METH_VARARGS, "Print out"},    
  {"gravity", gravity, METH_VARARGS, "Print out"},      
  {"manning_friction", manning_friction, METH_VARARGS, "Print out"},        
  //{"distribute_to_vertices_and_edges", 
  // distribute_to_vertices_and_edges, METH_VARARGS},    
  //{"update_conserved_quantities", 
  // update_conserved_quantities, METH_VARARGS},       
  //{"set_initialcondition", 
  // set_initialcondition, METH_VARARGS},    
  {NULL, NULL}
};
	
// Module initialisation   
void initshallow_water_ext(void){
  Py_InitModule("shallow_water_ext", MethodTable);
  
  import_array();     //Necessary for handling of NumPY structures  
}