#include "Python.h"
#include "numpy/arrayobject.h"
#include "math.h"
#include <stdio.h>
const double pi = 3.14159265358979;


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


/* double max(double a, double b) { */
/* 	double z; */
/* 	z=(a>b)?a:b; */
/* 	return z;} */
	
/* double min(double a, double b) { */
/* 	double z; */
/* 	z=(a<b)?a:b; */
/* 	return z;} */


// Function to obtain speed from momentum and depth.
// This is used by flux functions
// Input parameters uh and h may be modified by this function.
double _compute_speed(double *uh, 
		      double *h, 
		      double epsilon, 
		      double h0) {
  
  double u;
  
  if (*h < epsilon) {
    *h = 0.0;  //Could have been negative
     u = 0.0;
  } else {
    u = *uh/(*h + h0/ *h);    
  }
  

  // Adjust momentum to be consistent with speed
  *uh = u * *h;
  
  return u;
}

//-------------------------------------------------------------
// New vel based code
//-------------------------------------------------------------
//Innermost flux function (using w=z+h)
int _flux_function_vel(double *q_left, double *q_right,
		       double normals, double g, double epsilon, double h0,
		       double *edgeflux, double *max_speed) {
    
    int i;
    double flux_left[2], flux_right[2];
    double w_left, h_left, uh_left, z_left, u_left, soundspeed_left;
    double w_right, h_right, uh_right, z_right, u_right, soundspeed_right;
    double z, s_max, s_min, denom;
    

    w_left  = q_left[0];
    uh_left = q_left[1]*normals;
    z_left  = q_left[2];
    h_left  = q_left[3];
    u_left  = q_left[4]*normals;

/*     printf("w_left  = %f \n",w_left); */
/*     printf("uh_left = %f \n",uh_left); */
/*     printf("z_left  = %f \n",z_left); */
/*     printf("h_left  = %f \n",h_left); */
/*     printf("u_left  = %f \n",u_left); */
    
    w_right  = q_right[0];
    uh_right = q_right[1]*normals;
    z_right  = q_right[2];
    h_right  = q_right[3];
    u_right  = q_right[4]*normals;		
    
    z = (z_left+z_right)/2.0;		   
    
    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*u_left*h_left + 0.5*g*h_left*h_left;
    
    flux_right[0] = u_right*h_right;
    flux_right[1] = u_right*u_right*h_right + 0.5*g*h_right*h_right;
    
    // Flux computation
    denom = s_max-s_min;
    if (denom < epsilon) {
	for (i=0; i<2; i++) edgeflux[i] = 0.0;
	*max_speed = 0.0;
    } else {
	edgeflux[0] = s_max*flux_left[0] - s_min*flux_right[0];
	edgeflux[0] += s_max*s_min*(w_right-w_left);
	edgeflux[0] /= denom;
	edgeflux[1] = s_max*flux_left[1] - s_min*flux_right[1];
	edgeflux[1] += s_max*s_min*(uh_right-uh_left);
	edgeflux[1] /= denom;
	edgeflux[1] *= normals;
	
	// Maximal wavespeed
	*max_speed = max(fabs(s_max), fabs(s_min));
    }  
    return 0;		
}

// Computational function for flux computation
double _compute_fluxes_vel_ext(double cfl,
			       double timestep,
			       double epsilon,
			       double g,
			       double h0,
			       long* neighbours,
			       long* neighbour_vertices,
			       double* normals,
			       double* areas,
			       double* stage_edge_values,
			       double* xmom_edge_values,
			       double* bed_edge_values,
			       double* height_edge_values,
			       double* velocity_edge_values,
			       double* stage_boundary_values,
			       double* xmom_boundary_values,
			       double* bed_boundary_values,
			       double* height_boundary_values,
			       double* velocity_boundary_values,
			       double* stage_explicit_update,
			       double* xmom_explicit_update,
			       int number_of_elements,
			       double* max_speed_array) {
		
    double flux[2], ql[5], qr[5], edgeflux[2];
    double max_speed, normal;
    int k, i, ki, n, m, nm=0;

    //printf("Inside _comp\n");    
    
    for (k=0; k<number_of_elements; k++) {
	flux[0] = 0.0;
	flux[1] = 0.0;
	//printf("k = %d\n",k);

	for (i=0; i<2; i++) {
	    ki = k*2+i;
	    
	    ql[0] = stage_edge_values[ki];
	    ql[1] = xmom_edge_values[ki];
	    ql[2] = bed_edge_values[ki];
	    ql[3] = height_edge_values[ki];
	    ql[4] = velocity_edge_values[ki];
	    
	    n = neighbours[ki];
	    if (n<0) {
		m = -n-1;
		qr[0] = stage_boundary_values[m];
		qr[1] = xmom_boundary_values[m];
		qr[2] = bed_boundary_values[m];
		qr[3] = height_boundary_values[m];
		qr[4] = velocity_boundary_values[m];
	    } else {
		m = neighbour_vertices[ki];
		nm = n*2+m;
		qr[0] = stage_edge_values[nm];
		qr[1] = xmom_edge_values[nm];
		qr[2] = bed_edge_values[nm];
		qr[3] = height_edge_values[nm];
		qr[4] = velocity_edge_values[nm];
	    }
	    
	    normal = normals[ki];
	    _flux_function_vel(ql, qr, normal, g, epsilon, h0, edgeflux, &max_speed);
	    flux[0] -= edgeflux[0];
	    flux[1] -= edgeflux[1];
	    
	    // Update timestep based on edge i and possibly neighbour n
	    if (max_speed > epsilon) {
		// Original CFL calculation
		
		timestep = min(timestep, 0.5*cfl*areas[k]/max_speed); 
		if (n>=0) {
		    timestep = min(timestep, 0.5*cfl*areas[n]/max_speed); 
		}
	    }
	} // End edge i (and neighbour n)
	flux[0] /= areas[k];
	stage_explicit_update[k] = flux[0];
	flux[1] /= areas[k];
	xmom_explicit_update[k] = flux[1];
	
	//Keep track of maximal speeds
	max_speed_array[k]=max_speed;
    }
    return timestep;	
}

//-------------------------------------------------------------
// Old code
//------------------------------------------------------------
//Innermost flux function (using w=z+h)
int _flux_function(double *q_left, double *q_right,
        double z_left, double z_right,
		   double normals, double g, double epsilon, double h0,
		double *edgeflux, double *max_speed) {
		
		int i;
		double ql[2], qr[2], flux_left[2], flux_right[2];
		double z, w_left, h_left, uh_left, soundspeed_left, u_left;
		double w_right, h_right, uh_right, soundspeed_right, u_right;
		double s_max, s_min, denom;
		
		//printf("h0 = %f \n",h0);
		ql[0] = q_left[0];
		ql[1] = q_left[1];
		ql[1] = ql[1]*normals;
	
		qr[0] = q_right[0];
		qr[1] = q_right[1];
		qr[1] = qr[1]*normals;
	  
		z = (z_left+z_right)/2.0;		   
				   
		//w_left = ql[0];
		//h_left = w_left-z;
		//uh_left = ql[1];
		


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

		u_left = _compute_speed(&uh_left, &h_left, epsilon, h0);

		w_right = qr[0];
		h_right = w_right-z;
		uh_right = qr[1];

		u_right = _compute_speed(&uh_right, &h_right, epsilon, h0);
  
		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_right[0] = u_right*h_right;
		flux_right[1] = u_right*uh_right + 0.5*g*h_right*h_right;

		// Flux computation
		denom = s_max-s_min;
		if (denom < epsilon) {
			for (i=0; i<2; i++) edgeflux[i] = 0.0;
			*max_speed = 0.0;
		} else {
			edgeflux[0] = s_max*flux_left[0] - s_min*flux_right[0];
			edgeflux[0] += s_max*s_min*(qr[0]-ql[0]);
			edgeflux[0] /= denom;
			edgeflux[1] = s_max*flux_left[1] - s_min*flux_right[1];
			edgeflux[1] += s_max*s_min*(qr[1]-ql[1]);
			edgeflux[1] /= denom;
			edgeflux[1] *= normals;
    
		// Maximal wavespeed
        *max_speed = max(fabs(s_max), fabs(s_min));
		}  
    return 0;		
	}
		
		
	
	
// Computational function for flux computation
double _compute_fluxes_ext(
			   double cfl,
			   double timestep,
			   double epsilon,
			   double g,
			   double h0,
			   long* neighbours,
			   long* neighbour_vertices,
			   double* normals,
			   double* areas,
			   double* stage_edge_values,
			   double* xmom_edge_values,
			   double* bed_edge_values,
			   double* stage_boundary_values,
			   double* xmom_boundary_values,
			   double* stage_explicit_update,
			   double* xmom_explicit_update,
			   int number_of_elements,
			   double* max_speed_array) {
		
                double flux[2], ql[2], qr[2], edgeflux[2];
		double zl, zr, max_speed, normal;
		int k, i, ki, n, m, nm=0;
		
		
		for (k=0; k<number_of_elements; k++) {
			flux[0] = 0.0;
			flux[1] = 0.0;
			
			for (i=0; i<2; i++) {
				ki = k*2+i;
				
				ql[0] = stage_edge_values[ki];
				ql[1] = xmom_edge_values[ki];
				zl = bed_edge_values[ki];
				
				n = neighbours[ki];
				if (n<0) {
					m = -n-1;
					qr[0] = stage_boundary_values[m];
					qr[1] = xmom_boundary_values[m];
					zr = zl;
				} else {
					m = neighbour_vertices[ki];
					nm = n*2+m;
					qr[0] = stage_edge_values[nm];
					qr[1] = xmom_edge_values[nm];
					zr = bed_edge_values[nm];				
				}
				
				normal = normals[ki];
				_flux_function(ql, qr, zl, zr, normal, g, epsilon, h0, edgeflux, &max_speed);
				flux[0] -= edgeflux[0];
				flux[1] -= edgeflux[1];
				
				// Update timestep based on edge i and possibly neighbour n
				if (max_speed > epsilon) {
				    // Original CFL calculation
				    
				    timestep = min(timestep, 0.5*cfl*areas[k]/max_speed); 
				    if (n>=0) {
					timestep = min(timestep, 0.5*cfl*areas[n]/max_speed); 
				    }
				}
            } // End edge i (and neighbour n)
			flux[0] /= areas[k];
			stage_explicit_update[k] = flux[0];
			flux[1] /= areas[k];
			xmom_explicit_update[k] = flux[1];
			
			//Keep track of maximal speeds
			max_speed_array[k]=max_speed;
		}
		return timestep;	
	}

//=========================================================================
// Python Glue
//=========================================================================
PyObject *compute_fluxes_ext(PyObject *self, PyObject *args) {
  
   PyObject 
	*domain,
	*stage, 
	*xmom, 
	*bed;

    PyArrayObject 
	*neighbours, 
	*neighbour_vertices,
	*normals, 
	*areas,
	*stage_vertex_values,
	*xmom_vertex_values,
	*bed_vertex_values,
	*stage_boundary_values,
	*xmom_boundary_values,
	*stage_explicit_update,
	*xmom_explicit_update,
	*max_speed_array;
    
  double timestep, epsilon, g, h0, cfl;
  int number_of_elements;

    
  // Convert Python arguments to C
  if (!PyArg_ParseTuple(args, "dOOOO",
			&timestep,
			&domain,
			&stage,
			&xmom,
			&bed)) {
      PyErr_SetString(PyExc_RuntimeError, "comp_flux_vel_ext.c: compute_fluxes_ext could not parse input");
      return NULL;
  }


    epsilon           = get_python_double(domain,"epsilon");
    g                 = get_python_double(domain,"g");
    h0                = get_python_double(domain,"h0");
    cfl               = get_python_double(domain,"CFL");
 
    
    neighbours        = get_consecutive_array(domain, "neighbours");
    neighbour_vertices= get_consecutive_array(domain, "neighbour_vertices"); 
    normals           = get_consecutive_array(domain, "normals");
    areas             = get_consecutive_array(domain, "areas");    
    max_speed_array   = get_consecutive_array(domain, "max_speed_array");
    
    stage_vertex_values = get_consecutive_array(stage, "vertex_values");    
    xmom_vertex_values  = get_consecutive_array(xmom, "vertex_values");    
    bed_vertex_values   = get_consecutive_array(bed, "vertex_values");    

    stage_boundary_values = get_consecutive_array(stage, "boundary_values");    
    xmom_boundary_values  = get_consecutive_array(xmom, "boundary_values");    


    stage_explicit_update = get_consecutive_array(stage, "explicit_update");    
    xmom_explicit_update  = get_consecutive_array(xmom, "explicit_update");    



    number_of_elements = stage_vertex_values -> dimensions[0];


 
    // Call underlying flux computation routine and update 
    // the explicit update arrays 
    timestep = _compute_fluxes_ext(
	cfl,
	timestep,
	epsilon,
	g,
	h0,
	(long*) neighbours -> data,
	(long*) neighbour_vertices -> data,
	(double*) normals -> data,
	(double*) areas -> data,
	(double*) stage_vertex_values -> data,
	(double*) xmom_vertex_values -> data,
	(double*) bed_vertex_values -> data,
	(double*) stage_boundary_values -> data,
	(double*) xmom_boundary_values -> data,
	(double*) stage_explicit_update -> data,
	(double*) xmom_explicit_update -> data,
	number_of_elements,
	(double*) max_speed_array -> data);


  Py_DECREF(neighbours);
  Py_DECREF(neighbour_vertices);
  Py_DECREF(normals);
  Py_DECREF(areas);
  Py_DECREF(stage_vertex_values);
  Py_DECREF(xmom_vertex_values);
  Py_DECREF(bed_vertex_values);
  Py_DECREF(stage_boundary_values);
  Py_DECREF(xmom_boundary_values);
  Py_DECREF(stage_explicit_update);
  Py_DECREF(xmom_explicit_update);
  Py_DECREF(max_speed_array);




  // Return updated flux timestep
  return Py_BuildValue("d", timestep);
}


//------------------------------------------------
// New velocity based compute fluxes
//------------------------------------------------
PyObject *compute_fluxes_vel_ext(PyObject *self, PyObject *args) {
  
    PyObject 
	*domain,
	*stage, 
	*xmom, 
	*bed,
	*height,
	*velocity;

    PyArrayObject 
	*neighbours, 
	*neighbour_vertices,
	*normals, 
	*areas,
	*stage_vertex_values,
	*xmom_vertex_values,
	*bed_vertex_values,
	*height_vertex_values,
	*velocity_vertex_values,
	*stage_boundary_values,
	*xmom_boundary_values,
	*bed_boundary_values,
	*height_boundary_values,
	*velocity_boundary_values,
	*stage_explicit_update,
	*xmom_explicit_update,
	*max_speed_array;
    
  double timestep, epsilon, g, h0, cfl;
  int number_of_elements;

    
  // Convert Python arguments to C
  if (!PyArg_ParseTuple(args, "dOOOOOO",
			&timestep,
			&domain,
			&stage,
			&xmom,
			&bed,
			&height,
			&velocity)) {
      PyErr_SetString(PyExc_RuntimeError, "comp_flux_vel_ext.c: compute_fluxes_vel_ext could not parse input");
      printf("comp_flux_vel_ext.c: compute_fluxes_vel_ext could not parse input");
      return NULL;
  }


    epsilon           = get_python_double(domain,"epsilon");
    g                 = get_python_double(domain,"g");
    h0                = get_python_double(domain,"h0");
    cfl               = get_python_double(domain,"CFL");
 
    
    neighbours        = get_consecutive_array(domain, "neighbours");
    neighbour_vertices= get_consecutive_array(domain, "neighbour_vertices"); 
    normals           = get_consecutive_array(domain, "normals");
    areas             = get_consecutive_array(domain, "areas");    
    max_speed_array   = get_consecutive_array(domain, "max_speed_array");
    
    stage_vertex_values      = get_consecutive_array(stage,    "vertex_values");    
    xmom_vertex_values       = get_consecutive_array(xmom,     "vertex_values");    
    bed_vertex_values        = get_consecutive_array(bed,      "vertex_values");    
    height_vertex_values     = get_consecutive_array(height,   "vertex_values");    
    velocity_vertex_values   = get_consecutive_array(velocity, "vertex_values");    

    stage_boundary_values     = get_consecutive_array(stage,     "boundary_values");    
    xmom_boundary_values      = get_consecutive_array(xmom,      "boundary_values");    
    bed_boundary_values       = get_consecutive_array(bed,       "boundary_values");    
    height_boundary_values    = get_consecutive_array(height,    "boundary_values");    
    velocity_boundary_values  = get_consecutive_array(velocity,  "boundary_values");    


    stage_explicit_update = get_consecutive_array(stage, "explicit_update");    
    xmom_explicit_update  = get_consecutive_array(xmom,  "explicit_update");    

    number_of_elements = stage_vertex_values -> dimensions[0];
 

    //printf("comp_flux_vel 1 N = %d %15.5e %15.5e %15.5e %15.5e %15.5e \n",number_of_elements,cfl,timestep,epsilon,g,h0);



    // Call underlying flux computation routine and update 
    // the explicit update arrays 
    timestep = _compute_fluxes_vel_ext(cfl,
				       timestep,
				       epsilon,
				       g,
				       h0,
				       (long*) neighbours -> data,
				       (long*) neighbour_vertices -> data,
				       (double*) normals -> data,
				       (double*) areas -> data,
				       (double*) stage_vertex_values -> data,
				       (double*) xmom_vertex_values -> data,
				       (double*) bed_vertex_values -> data,
				       (double*) height_vertex_values -> data,
				       (double*) velocity_vertex_values -> data,
				       (double*) stage_boundary_values -> data,
				       (double*) xmom_boundary_values -> data,
				       (double*) bed_boundary_values -> data,
				       (double*) height_boundary_values -> data,
				       (double*) velocity_boundary_values -> data,
				       (double*) stage_explicit_update -> data,
				       (double*) xmom_explicit_update -> data,
				       number_of_elements,
				       (double*) max_speed_array -> data);

    //printf("comp_flux_vel 2 \n");    
    
    Py_DECREF(neighbours);
    Py_DECREF(neighbour_vertices);
    Py_DECREF(normals);
    Py_DECREF(areas);
    Py_DECREF(stage_vertex_values);
    Py_DECREF(xmom_vertex_values);
    Py_DECREF(bed_vertex_values);
    Py_DECREF(height_vertex_values);
    Py_DECREF(velocity_vertex_values);
    Py_DECREF(stage_boundary_values);
    Py_DECREF(xmom_boundary_values);
    Py_DECREF(bed_boundary_values);
    Py_DECREF(height_boundary_values);
    Py_DECREF(velocity_boundary_values);
    Py_DECREF(stage_explicit_update);
    Py_DECREF(xmom_explicit_update);
    Py_DECREF(max_speed_array);
    
    //printf("comp_flux_vel 3 \n");

    // Return updated flux timestep
    return Py_BuildValue("d", timestep);
}



//-------------------------------
// Method table for python module
//-------------------------------

static struct PyMethodDef MethodTable[] = {
  {"compute_fluxes_ext", compute_fluxes_ext, METH_VARARGS, "Print out"},
  {"compute_fluxes_vel_ext", compute_fluxes_vel_ext, METH_VARARGS, "Print out"},
  {NULL, NULL}
};

// Module initialisation
void initsww_vel_comp_flux_ext(void){
  Py_InitModule("sww_vel_comp_flux_ext", MethodTable);
  import_array();
}