#include <stdlib.h>
#include "cv.h"
#include "leastsquare.h"

float summation(int size, float* data) {
  float total = 0.0;
  int i;
  for (i = 0; i < size; i++) {
    total += data[i];
  }
  return total;
}

float expt(float base, int pow) {
  float ans = 1.0;
  while (pow > 0) {
    ans = ans * base;
    pow --;
  }
  return ans;
}

void compute_least_square(int size, float* data_xx, float* data_yy,
			  float* a, float* b, float* c, float* d) {

  int i;
  float x6, x5, x4, x3, x2, x1, x0;
  float y3, y2, y1, y0;
  CvMat A;
  CvMat Ainv;

  /* compute all the x's and y's */
  x6 = x5 = x4 = x3 = x2 = x1 = x0 = 0.0;
  y3 = y2 = y1 = y0 = 0.0;
  for (i = 0; i < size; i++) {
    x6 += expt(data_xx[i],6);
    x5 += expt(data_xx[i],5);
    x4 += expt(data_xx[i],4);
    x3 += expt(data_xx[i],3);
    x2 += expt(data_xx[i],2);
    x1 += data_xx[i];
    x0 += 1.0;

    y3 += expt(data_xx[i],3) * data_yy[i];
    y2 += expt(data_xx[i],2) * data_yy[i];
    y1 += data_xx[i] * data_yy[i];
    y0 += data_yy[i];
  }

  /* construct A */
  A = cvMat(4, 4, CV_32F, NULL);
  cvmAlloc(&A);
  cvmSet(&A, 0, 0, x6);
  cvmSet(&A, 0, 1, x5);
  cvmSet(&A, 0, 2, x4);
  cvmSet(&A, 0, 3, x3);

  cvmSet(&A, 1, 0, x5);
  cvmSet(&A, 1, 1, x4);
  cvmSet(&A, 1, 2, x3);
  cvmSet(&A, 1, 3, x2);

  cvmSet(&A, 2, 0, x4);
  cvmSet(&A, 2, 1, x3);
  cvmSet(&A, 2, 2, x2);
  cvmSet(&A, 2, 3, x1);

  cvmSet(&A, 3, 0, x3);
  cvmSet(&A, 3, 1, x2);
  cvmSet(&A, 3, 2, x1);
  cvmSet(&A, 3, 3, x0);

  /* compute A^-1 */
  Ainv = cvMat(4, 4, CV_32F, NULL);
  cvmAlloc(&Ainv);
  cvmInvert(&A, &Ainv);

  /* x = A^-1*y */
  *a =
    cvmGet(&Ainv,0,0)*y3 +
    cvmGet(&Ainv,0,1)*y2 +
    cvmGet(&Ainv,0,2)*y1 +
    cvmGet(&Ainv,0,3)*y0;
  *b =
    cvmGet(&Ainv,1,0)*y3 +
    cvmGet(&Ainv,1,1)*y2 +
    cvmGet(&Ainv,1,2)*y1 +
    cvmGet(&Ainv,1,3)*y0;
  *c =
    cvmGet(&Ainv,2,0)*y3 +
    cvmGet(&Ainv,2,1)*y2 +
    cvmGet(&Ainv,2,2)*y1 +
    cvmGet(&Ainv,2,3)*y0;
  *d =
    cvmGet(&Ainv,3,0)*y3 +
    cvmGet(&Ainv,3,1)*y2 +
    cvmGet(&Ainv,3,2)*y1 +
    cvmGet(&Ainv,3,3)*y0;

  cvmFree(&A);
  cvmFree(&Ainv);
}


/* matlab code used for testing

data_xx = [1 2 4 5 6 7];
data_yy = [10 6 3 4 6 10];

figure(1); clf; hold on;
plot(data_xx, data_yy, 'x:');

% compute sum(x^4), sum(x^3), sum(x^2), sum(x^1), sum(x^0)
x4 = sum(data_xx .^ 4);
x3 = sum(data_xx .^ 3);
x2 = sum(data_xx .^ 2);
x1 = sum(data_xx .^ 1);
x0 = sum(data_xx .^ 0);

% compute sum(x^2*y), sum(x*y), sum(y)
y2 = sum(data_xx .^ 2 .* data_yy);
y1 = sum(data_xx .* data_yy);
y0 = sum(data_yy);

% construct A and y
A = [ x4 x3 x2;  x3 x2 x1;  x2 x1 x0 ];
y = [ y2;  y1;  y0];

% solve for least square solution with a ax^2 + bx + c = y
x = A^-1 * y;
a = x(1);
b = x(2);
c = x(3);


eqn_xx = min(data_xx):0.1:max(data_xx);
eqn_yy = (a * eqn_xx .^ 2) + (b * eqn_xx) + c;
plot(eqn_xx, eqn_yy, 'r');

*/
