/*  locpoly.c                          */
/*  version 1.0.1                      */
/*  Roberto G. Gutierrez, StataCorp    */

#include <math.h>
#include <stdio.h>
#include <string.h>
#include <malloc.h>

#define  SQRTTWOPI   2.50662827463100024
#define  PI          3.14159265358979312

#define  ABS(a)      ( (a>0) ? a : -(a) )

#include "stplugin.h"

/* Forward references */
static double 		biweight(double x) ;
static double 		cosine(double x) ;
static double 		triangle(double x) ;
static double 		parzen(double x) ;
static double 		gaussian(double x) ;
static double 		rectangle(double x) ;
static double 		epanechnikov(double x) ;
static double		*dvector(int n) ;
static void		free_dvector(double *x) ;
static double		**dmatrix(int n, int m) ;
static void		free_dmatrix(double **x, int n) ;
static int   		invert_dmatrix(double **A, int n) ;
static ST_retcode	syntaxerror() ;
static ST_retcode 	locpoly_wrk(ST_int ng, ST_double h, ST_int p, 
				    double (*K)(double)) ;

STDLL stata_call(ST_int argc, char *argv[])
{
	ST_double	width ;
	ST_int		p, ngrid ;
	double		(*kernel)(double) ;

	if(SF_nvars() != 4) {                    /* number of vars */
		return(syntaxerror()) ;
	}

	if(argc != 4) { 			 /* number of arguments */
		return(syntaxerror()) ;
	}

	if(sscanf(argv[0], "%d", &ngrid) != 1) {     /* grid size */
		return(syntaxerror()) ;
	}

	if(sscanf(argv[1], "%lf", &width) != 1) {    /* bandwidth */
		return(syntaxerror()) ;
	}

	if(sscanf(argv[2], "%d", &p) != 1) {         /* degree */
		return(syntaxerror()) ;
	}
						  /* match kernel to func */
	if(strcmp(argv[3],"biweight")==0) kernel = biweight ;
	else if(strcmp(argv[3],"cosine")==0) kernel = cosine ;
	else if(strcmp(argv[3],"triangle")==0) kernel = triangle ;
	else if(strcmp(argv[3],"parzen")==0) kernel = parzen ;
	else if(strcmp(argv[3],"gaussian")==0) kernel = gaussian ;
	else if(strcmp(argv[3],"rectangle")==0) kernel = rectangle ;
	else if(strcmp(argv[3],"epanechnikov")==0) kernel = epanechnikov ;
	else return(syntaxerror()) ;

	return(locpoly_wrk((int) ngrid, (double) width, (int) p, kernel)) ; 
}

static ST_retcode syntaxerror() 
{
	SF_error("syntax error\n") ;
	SF_error("Usage: _lpwork y x xgrid yhat, ") ;
	SF_error("n width degree kernel\n") ;

	return((ST_retcode) 198) ;
}

static ST_retcode locpoly_wrk(int ng, double h, int p, 
			      double (*K)(double))
{
	double **A, *B, *C ;
	double arg, Kh, prod;
	double z1, z2 ,sum ;
	int y, x, xg, yhat ;
	int i, j, k, ii ;
	int p1 = p + 1;
	
	y = 1; x = 2; xg = 3; yhat = 4 ;  /* positional variables */
	
	A = dmatrix(p1, p1) ;
	B = dvector(p1) ;
	C = dvector(2*p + 1) ;

	if (!A || !B || !C) return((ST_retcode) 909) ;  /* mem alloc error */

	for(i=1;i<=(int) ng;i++) {
		if (SF_poll()) return(SW_stopflag) ;	
		for(k=1;k<=(2*p+1);k++) C[k]=0.0 ;
		for(k=1;k<=p1;k++) B[k]=0.0 ;
		for(j=1;j<=SF_nobs();j++) {
			if(SF_ifobs(j)) {
				(void) SF_vdata(x,j,&z1) ; 
				(void) SF_vdata(xg,i,&z2) ;
				arg = (z1-z2) ;
				Kh = (*K)(arg/h) ;
				for(k=1,prod=1.0; k<=(2*p+1); k++) {
					C[k] += Kh*prod ;	
					prod *= arg ;
				}
				for(k=1,prod=1.0; k<=p1; k++) {
					(void) SF_vdata(y,j,&z1) ;
					B[k] += Kh*z1*prod ;
					prod *= arg ;
				}
			}
		}

		for(k=1; k<=p1; k++) 
			for(ii=1; ii<=p1; ii++) 
				A[k][ii] = C[k+ii-1] ;

		if(invert_dmatrix(A,p1)!=0) {
			for(k=1, sum = 0.0; k<=p1; k++) {
				sum += A[1][k]*B[k] ;
			}
			(void) SF_vstore(yhat,i,sum) ;
		}
	}

  	free_dmatrix(A,p1);
  	free_dvector(B);
  	free_dvector(C);

	return((ST_retcode) 0) ;
}

static double biweight(double x)
{
	return((ABS(x)<1.0) ? 0.9375*(1.0-x*x)*(1.0-x*x) : 0.0);
}

static double cosine(double x)
{
  	return((ABS(x)<0.5) ? 1 + cos(2.0*PI*x) : 0.0);
}

static double triangle(double x)
{
  	return((ABS(x)<1.0) ? 1.0-ABS(x) : 0.0);
}

static double parzen(double x) 
{
  	if(ABS(x)<=0.5) 
    		return(4.0/3.0 - 8.0*x*x + 8*ABS(x)*ABS(x)*ABS(x));
  	else if(ABS(x)<1.0)
    		return(8.0/3.0*(1.0-ABS(x))*(1.0-ABS(x))*(1.0*ABS(x)));
  	else
    		return(0.0) ;
}  

static double gaussian(double x)
{
  	return(exp(-x*x/2.0)/SQRTTWOPI);
}

static double rectangle(double x)
{
  	return((ABS(x)<1.0) ? 0.5 : 0.0);
}

static double epanechnikov(double x)
{
  	return((ABS(x)<sqrt(5.0)) ? 3.0/(4.0*sqrt(5.0))*(1.0-x*x/5.0) : 0.0);
}

static double *dvector(int n)
{
   	double *v ;

   	v = (double *) malloc((unsigned) n*sizeof(double)) ;
   	if(!v) return(NULL) ;
  	return(v-1) ;
}

static void free_dvector(double *x)
{
   	free((char*) (x+1)) ;
}

static double **dmatrix(int n,int m)
{
   	int i ;
   	double **x ;

   	x=(double **) malloc((unsigned) n*sizeof(double)) ;
   	if(!x) return(NULL) ;
   	x-=1;

   	for(i=1;i<=n;i++) {
        	x[i]=(double *) malloc((unsigned) m*sizeof(double)) ;
      		if(!x[i]) return(NULL) ;
      		x[i]-=1 ;
   	}
    
   	return(x) ;
}  

static void free_dmatrix(double **x,int n)
{
   	int i;

   	for(i=n;i>=1;i--) free((char*) (x[i]+1));
   	free((char*) (x+1));
}

static int invert_dmatrix(double **A,int n)
{
  	int i,j,k;
  	double c,eps=1.e-8;
  	for(k=1;k<=n;k++) {
    		if(ABS(c=A[k][k])<eps) return(0);   
    		for(i=1;i<=n;i++) 
      			for(j=1;j<=n;j++) 
				if(i!=k && j!=k) A[i][j]-=(A[i][k]*A[k][j])/c;
		for(i=1;i<=n;i++) { 
	  		A[i][k]/=-c; 
	  		A[k][i]/=c;
		}
		A[k][k]=1/c;
  	}
  	return(1);
}