// JAMA is a semi-standard matrix libraries in Java:
import Jama.util.*;
import Jama.*;

// a bunch of different radial basis functions to test out
int RBF_TPS = 0;
int RBF_DCUBED = 1;
int RBF_LIN = 2;
int RBF_FIFTH = 3;
int RBF_COUNT = 4;
int RBF_DSQ = 4;
int RBF_FOURTH = 5;
int RBF_WENDLAND_C2_3D = 6;
int RBF_CNT = 7;
String g_fnNames[] = {"(r^2)*log(r) (\"thin plate\")", "r^3", "r", "r^5", "r^2", "r^4", "wendland"};

// we can add some extra polynomial terms to improve the solution; here I just added up to linear
int POLY_NONE = 0;
int POLY_CONST = 1;
int POLY_LIN = 2;
int POLY_COUNT = 3;
int g_extraParams[] = {0,1,3}; 
String g_polyNames[] = {"zero", "constant", "linear"};


class rbf {
  int rbfFn = RBF_TPS; // which rbf we're using right now
  int polyMode = 2; // which polynomial we're using
  
  Matrix paramsMat; // stores the current weights for all the radial basis functions
  ArrayList draggables; // our control points are class Draggable (which gives some UI to them)
  
  rbf() {
    draggables = new ArrayList();
  }
  
  // add a control point
  void add(int x, int y) {
    float z = eval(x, y);
    Draggable d = new Draggable(x, y, 1, 5);
    d.z = z;
    draggables.add(d);
  }

  // big collection of possible radial basis functions 
  float evalRBF(float x, float y, float ox, float oy) {
    float dx = x-ox;
    float dy = y-oy;
    float dsq = dx*dx + dy*dy;
    float d = sqrt(dsq);
    if (rbfFn == RBF_TPS)
      return dsq * log(d+.1);
    else if (rbfFn == RBF_DCUBED)
      return d*d*d;
    else if (rbfFn == RBF_DSQ)
      return dsq;
    else if (rbfFn == RBF_LIN)
      return d;
    else if (rbfFn == RBF_FOURTH)
      return d*d*d*d;
    else if (rbfFn == RBF_FIFTH)
      return d*d*d*d*d;
    else if (rbfFn == RBF_WENDLAND_C2_3D) {
      float r = d/200.0; // 100 determines the range of influence of the control point -- should really be exposed as a parameter
      float omr = 1-r; // if this would be negative, we return zero instead
      if (r < 1) {
        return omr*omr*omr*omr*(4*r+1);
      } else {
        return 0.0;
      }
    }
    return 0.0;
  }

  // set of helper functions for calling the above radial basis function:
  // evaluate what the rbf function at control point i contributes to the location of control point j (or vice-versa, it's symmetric) (before any scaling)
  float evalRBF(int i, int j) {
    Draggable a = (Draggable)draggables.get(i);
    Draggable b = (Draggable)draggables.get(j);
    return evalRBF(a.x,a.y,b.x,b.y);
  }
  // evaluate what the rbf function at control point i contributes to the location x,y
  float evalRBF(int i, float x, float y) {
    Draggable a = (Draggable)draggables.get(i);
    return evalRBF(a.x,a.y,x,y);
  }
  // eval the full sum at a given position and for given weights 
  float evalRBF(Matrix params, float x, float y) {
    float sum = 0;
    int dsize = draggables.size();
    for (int i = 0; i < dsize; i++) {
      sum += (float)params.get(i,0) * evalRBF(i,x,y);
    }
    if (polyMode > POLY_NONE) {
      sum += (float)(params.get(dsize,0));
      if (polyMode > POLY_CONST) {
        sum += (float)(x*params.get(dsize+1,0) + y*params.get(dsize+2,0));
      }
    }
    return sum;
  }
  // simplest helper function; evaluates full sum at given position
  float eval(float x, float y) {
    if (draggables.size() == 1) {
      Draggable d = (Draggable)draggables.get(0);
      return d.z;
    }
    return evalRBF(paramsMat, x, y);
  }
  
  // move control points in response to mouse input, then re-solve for our weights
  void update() {
    for (int i = 0; i < draggables.size(); i++) {
      ((Draggable)draggables.get(i)).update(10, mouseX - pmouseX, mouseY - pmouseY);
    }
    updateParams();
  }
  

  
  // returns solution matrix of parameters, which is really a vector since it has just 1 col with dsize+[number of polynomial terms] rows
  Matrix solveForRBFParams() {
    int dsize = draggables.size(); // dsize == number of radial basis functions
    int extra = g_extraParams[polyMode]; // extra == number of polynomial terms
    Matrix M = new Matrix(dsize+extra,dsize+extra); // left-hand side matrix as in M x = B
    Matrix B = new Matrix(dsize+extra,1); // right hand side 
    for (int i = 0; i < dsize; i++) {
      // fill in the upper-left square of rbf values
      Draggable d = (Draggable)draggables.get(i);
      M.set(i,i,evalRBF(i,i)); 
      for (int j = i+1; j < dsize; j++) {
        double rij = evalRBF(i,j);
        M.set(i,j,rij);
        M.set(j,i,rij);
      }
      // fill in the bottom and right rectangles with polynomial values
      if (polyMode > POLY_NONE) {
        // add const term of polynomial
        M.set(dsize,i,1);
        M.set(i,dsize,1);
        if (polyMode > POLY_CONST) {
          // add linear term of polynomial
          M.set(dsize+1,i,d.x);
          M.set(i,dsize+1,d.x);
          M.set(dsize+2,i,d.y);
          M.set(i,dsize+2,d.y);
        }
      }
      // fill in the rhs with target values
      B.set(i,0,d.z);
    }
    LUDecomposition lud = new LUDecomposition(M); // perform LU factorization as first step in solving the matrix
    
    if (lud.isNonsingular()) { // if solution exists, return it
      return lud.solve(B);
    } else {
      println("matrix is singular!");
      return new Matrix(dsize+extra,1); // everything is zero ...
    }
  }
  // simplest helper
  void updateParams() {
     paramsMat = solveForRBFParams();
  }
  
  
  // display the control point UI stuff
  void display() {
    for (int i = 0; i < draggables.size(); i++) {
      ((Draggable)draggables.get(i)).display();
    }
  }
  // delete a control point
  void removeUnderMouse() {
    for (int i = 0; i < draggables.size(); i++) {
      Draggable d = (Draggable)draggables.get(i);
      if (d.over) {
        draggables.remove(i);
        return;
      }
    }    
  }



}