#include <math.h>
#include <nr.h>
#include <nrutil.h>
#include <tcl.h>

static void (*derivs)() = NULL;
static float *param = NULL;
static int nparam = 0;
static int ndimen = 0;
static float *y = NULL;
static float *dydx = NULL;

static void deriv_pendulum(x, y, dydx)
     float x;
     float *y;
     float *dydx;
{
  dydx[1] = y[2];
  dydx[2] = -param[1]*sin(y[1]) - param[2] * y[2];
}

typedef struct _FunctionRec {
  char *name;
  int ndimen;
  int nparam;
  void (*function)();
} Function;

Function function_list[] = {
  {"pendulum", 2, 2, deriv_pendulum},
  {NULL, 0, 0, NULL}
};

static int OdeValueCmd (dummy, interp, argc, argv)
     ClientData dummy;
     Tcl_Interp *interp;
     int argc;
     char **argv;
{
  double x;
  char buffer[256];
  int i;

  if (argc <= ndimen+1) {
    Tcl_SetResult(interp, "Not enough direction info", TCL_STATIC);
    return TCL_ERROR;
  }
  x = atof(argv[1]);
  for (i=1; i<=ndimen; i++)
    y[i] = atof(argv[i+1]);
  if (derivs)
    (*derivs)(x, y, dydx);
  for (i=1; i<= ndimen; i++) {
    sprintf(buffer, "%f", dydx[i]);
    Tcl_AppendElement(interp, buffer, 0);
  }
  return (TCL_OK);
}


static int OdeDirectionCmd (dummy, interp, argc, argv)
     ClientData dummy;
     Tcl_Interp *interp;
     int argc;
     char **argv;
{
  double x, sum;
  char buffer[256];
  int i;

  if (argc <= ndimen+1) {
    Tcl_SetResult(interp, "Not enough direction info", TCL_STATIC);
    return TCL_ERROR;
  }
  x = atof(argv[1]);
  for (i=1; i<=ndimen; i++)
    y[i] = atof(argv[i+1]);
  if (derivs)
    (*derivs)(x, y, dydx);

  sum = 0.0;
  for (i=1; i<=ndimen; i++) {
    sum += dydx[i]*dydx[i];
  }
  sum = sqrt(sum);

  if (sum == 0.0) {
    sum = 1.0;
  }

  for (i=1; i<= ndimen; i++) {
    sprintf(buffer, "%f", dydx[i] / sum);
    Tcl_AppendElement(interp, buffer, 0);
  }
  return (TCL_OK);
}

static int OdeSetFunctionCmd(dummy, interp, argc, argv) 
     ClientData dummy;
     Tcl_Interp *interp;
     int argc;
     char **argv;
{
  int i;
  for(i=0; function_list[i].name != NULL; i++) 
    if (!strcmp(argv[1], function_list[i].name)) {
      if (nparam)
	free_vector(param, 1, nparam);
      if (ndimen) {
	free_vector(y, 1, ndimen);
	free_vector(dydx, 1, ndimen);
      }
      nparam = function_list[i].nparam;
      ndimen = function_list[i].ndimen;
      derivs = function_list[i].function;
      
      param = vector(1, nparam);
      y = vector(1, ndimen);
      dydx = vector(1, ndimen);
      return (TCL_OK);
    }
  Tcl_SetResult(interp, "function not found", TCL_STATIC);
  return (TCL_ERROR);
}

/*
argv[1] = x value
argv[2] = y value
argv[3] = step size
argv[4] = number steps
*/

static int OdeIntegrateCmd(dummy, interp, argc, argv)
     ClientData dummy;
     Tcl_Interp *interp;
     int argc;
     char **argv;
{
  float x = atof(argv[1]);
  float step = atof(argv[3]);
  int nstep = atoi(argv[4]);
  int i, j;
  char buffer[256];

  int ninitial;
  char **initials;
  Tcl_SplitList(interp, argv[2], &ninitial, &initials);

  if (ninitial < ndimen) {
    for (i=0; i < ninitial; i++)
      free(initials[i]);
    Tcl_SetResult(interp, "Not enough initial conditions", TCL_STATIC);
    return TCL_ERROR;
  }

  for (i=0; i< ndimen; i++) 
    y[i+1] = atof(initials[i]);


  for (i=0; i < ninitial; i++)
    free(initials[i]);
  
  for (i=0; i< nstep; i++) {
    (*derivs)(x, y, dydx); 
    rk4(y, dydx, ndimen, x, step, y, derivs);
    for (j=1; j<=ndimen; j++) {
      sprintf(buffer, "%f", y[j]);
      Tcl_AppendElement(interp, buffer, 0);
    }
  }
  return (TCL_OK);
}

static int OdeSetParamsCmd(dummy, interp, argc, argv)
     ClientData dummy;
     Tcl_Interp *interp;
     int argc;
     char **argv;
{
  int i;
  if (argc < nparam + 1) 
    return (TCL_ERROR);

  for (i=1; i<= nparam; i++) 
    param[i] = atof(argv[i]);
  
  return (TCL_OK);
}

void init_ode (interp) 
     Tcl_Interp *interp;
{
  Tcl_CreateCommand(interp, "OdeDirection", OdeDirectionCmd,
		    (ClientData) NULL, (void (*)()) NULL);
  Tcl_CreateCommand(interp, "OdeIntegrate", OdeIntegrateCmd,
		    (ClientData) NULL, (void (*)()) NULL);
  Tcl_CreateCommand(interp, "OdeSetFunction", OdeSetFunctionCmd,
		    (ClientData) NULL, (void (*)()) NULL);
  Tcl_CreateCommand(interp, "OdeSetParams", OdeSetParamsCmd,
		    (ClientData) NULL, (void (*)()) NULL);
  Tcl_CreateCommand(interp, "OdeValue", OdeValueCmd,
		    (ClientData) NULL, (void (*)()) NULL);
}