/* ----------------------------------------------------------------------------

unireg.cc

Kernel regression estimator.  Self documented, execute unireg with no arguments
for documentation.

Copyright (C) 1991, 1992, 1993, 1994, 2002, 2006.

A. Ronald Gallant
Post Office Box 659 
Chapel Hill NC 27514-0659 
USA   

Permission to use, copy, modify, and distribute this software and its docu-
mentation for any purpose and without fee is hereby granted, provided that 
the above copyright notice appears in all copies and that both that copyright 
notice and this permission notice appear in supporting documentation.  This 
software is provided "as is" without any expressed or implied warranty.

-----------------------------------------------------------------------------*/

#include "libscl.h"
using namespace std;
using namespace scl;

INTEGER max_sample_size=1024;   // Can be changed with flag -m
INTEGER grid=100;               // Can be changed with flag -g
REAL    bandwidth;              // Set with flag -b, defaults otherwise.
REAL    xmin;                   // Set with flag -xmin, defaults otherwise.
REAL    xmax;                   // Set with flag -xmax, defaults otherwise.
int     predict = 0;            // The flag -p sets to true, estimate otw.
int     eformat = 0;            // The flag -e sets to true
int     cross_validate = 0;     // The flag -c sets to true, estimate otw.
int     normal_kernel = 1;      // The flag -q sets to false.
int     quadratic_kernel = 0;   // The flag -q sets to true.
int     filter = 0;             // The flag -f sets to true.
int     do_not_bin = 0;         // The flag -d sets to true.
int     compute_bandwidth = 1;
int     compute_xmax = 1;              
int     compute_xmin = 1;

INTEGER  n, nbin;
REAL    *Ysrt,     *Xsrt;
REAL    *Ysrt_top, *Xsrt_top;
REAL    *Ybin,     *Xbin;
REAL    *Ybin_top, *Xbin_top;
REAL    *Yvar,     *Xwgt;
REAL     xmean,xsigma,xspread;

void process_file(istream& cfile);
void usage_msg();
void estimate();
void pred();
void bin(REAL b);
REAL crossval(int isw);

REAL const zero = (REAL) 0.0;

char    *progname;

void usage_msg()
{
  cerr << '\n'
    << "usage: " <<progname<< " [options] filename \n\n"
    << "options: -p             predict at data points only \n"
    << "         -e             output observations in e format with full precision\n"
    << "         -c             cross validate instead of estimate \n" 
    << "         -q             quadratic kernel, normal is default\n"
    << "         -d             do not bin the data\n"
    << "         -f             acts as a filter, takes standard input\n"
    << "         -m     value   maximum sample size \n" 
    << "         -b     value   bandwidth\n"
    << "         -g     value   number of plot points sent to cout\n"
    << "         -xmin  value   minimum plot point\n"
    << "         -xmax  value   maximum plot point\n\n"
    << "example: "<<progname<<" -b 6.0 -q -g 200 in.dat >out.dat\n\n"
    << "The input file contains y x pairs, one pair per line, y first.\n"
    << "The output file is the same.\n\n";
}

int main(int argc, char *argp[], char *envp[])
{
  int      argx, flagx, filex;
  char     ch, flag;
  REAL     cv, cv_min, bandwidth_min, decrement, tmp, A, silverman;

  ch = envp[0][0]; ch++; // This supresses a compiler warning.

  progname = argp[0];

  // begin: loop over flags and input files.

  argx = 0;  filex = 1;  

  while (argx < argc) {

    if (argx==0) argx++;

    // begin: read flags, if any.

    while (argx < argc && (argp[argx])[0]=='-') {

      // Read flags, if any.  Note that flags can occur singly as -a -b 
      // or they can be grouped as -ab.  Just insert code within the
      // appropriate case and don't worry about it.

      for (flagx=1; (flag=(argp[argx])[flagx])!=0; flagx++) {

        // If a flag can have more than one letter, deal with it within
        // the appropriate case by stripping off letters and incrementing
        // flagx.

        // If a flag indicates that argp[argx+1] is a file that should be read 
        // then, within the appropriate case, open arg[argx+1] or otherwise 
        // deal with the file name arg[argx+1], assign argp[argx] to 
        // argp[argx+1], and increment argx.  

        // If a flag indicates that argp[argx+1] is a string containing the  
        // value of a parameter, etc., then, within the appropriate case, deal 
        // with the string arg[argx+1], assign argp[argx] to argp[argx+1], and 
        // increment argx.  

        switch (flag) {

        case 'c' :
          cross_validate = 1;
          break;

        case 'p' :
          predict = 1;
          break;

        case 'e' :
          eformat = 1;
          break;

        case 'd' :
          do_not_bin = 1;
          break;

        case 'q' :
          normal_kernel = 0;
          quadratic_kernel = 1;
          break;

        case 'f' :
          filter = 1;
          break;

        case 'b' :
          bandwidth = (REAL) atof(argp[argx+1]); 
          compute_bandwidth = 0;
          argp[argx+1] = argp[argx];
          argx++;
          break;

        case 'm' :
          max_sample_size = (INTEGER) atoi(argp[argx+1]); 
          argp[argx+1] = argp[argx];
          argx++;
          break;

        case 'g' :
          grid = (INTEGER) atoi(argp[argx+1]); 
          argp[argx+1] = argp[argx];
          argx++;
          break;

        case 'x' : 
          if (memcmp(&(argp[argx][flagx]),"xmin",4)==0) {
            flagx += 3;
            xmin = atof(argp[argx+1]); 
            argp[argx+1] = argp[argx];
            compute_xmin = 0;
            argx++;
            break;
          }
          if (memcmp(&(argp[argx][flagx]),"xmax",4)==0) {
            flagx += 3;
            xmax = atof(argp[argx+1]); 
            argp[argx+1] = argp[argx];
            compute_xmax = 0;
            argx++;
            break;
          }
          goto badflag;

        badflag :
        default :
          cerr << "Flag was ";
          while ((flag=(argp[argx])[flagx++])!=0) 
            cerr << flag;
          cerr << "\nError: unknown flag.\n";
          exit(1);
          break;

        } // end: switch

      } // end: for 

    argx++;

    } // end: while

    // end: read flags, if any.

    if (filter) {
      process_file(cin);
    }
    else {
      if (argx >= argc) {
        usage_msg();
        exit(1);
      }

      ifstream infile(argp[argx]);
      if (!infile) {
        cerr << "Error, " << progname << ", Cannot open " << argp[argx] << "\n";
        usage_msg();
        exit(1);
      } 
      istream& from = infile;

      process_file(from);
      infile.close();
    }

    if (compute_bandwidth) {
      A = (xsigma < xspread) ? xsigma : xspread;
      silverman = bandwidth = 0.9*A/pow((double)n,0.2);

      decrement = bandwidth/10.0;
      bandwidth = silverman + 5.0*decrement;
      bandwidth_min = silverman - 7.001*decrement;

      bin(bandwidth_min/3.1);

      tmp=bandwidth;
      cv_min = crossval(0);
      cerr << "\nTrials:\n";
      cerr << "cv " << cv_min << "  bw " << bandwidth << "\n";

      do{
        bandwidth -= decrement; 
        cv = crossval(0);
        if (fabs(bandwidth-silverman)<0.25*decrement) 
          cerr<<"cv "<<cv<<"  bw "<<bandwidth
              <<" = .9*min(sd,iqr/1.34)/pow(n,0.2)\n";
        else
          cerr << "cv " << cv << "  bw " << bandwidth << "\n";
        if (cv < cv_min) { 
          cv_min = cv; 
          tmp = bandwidth;
        }
      } while (bandwidth > bandwidth_min);
      bandwidth = tmp;
      cerr << "\nSelection:\n";
      cerr << "bandwidth = " << bandwidth << "\n\n";
    }
    else {
      bin(bandwidth/3.1);
    }

    if (cross_validate) {
      cv = crossval(1);
      exit(0);
    }

    if (predict) {
      pred();
      exit(0);
    }

    estimate();
    exit(0);

    filex++;
    argx++;
  } 

  // end: loop over flags and input files.

  exit(0);

}

void process_file(istream& from)
{ 
  REAL     *x;
  INTEGER  lo, hi, i;
  void     usage_msg();

  REAL     *Y,     *X;
  REAL     *Y_top, *X_top;

  Y = new REAL[max_sample_size+1];
  X = new REAL[max_sample_size+1];

  n = 0;
  Y_top=Y; X_top=X;

  while( (from>>*Y_top) && (from>>*X_top) && (n <= max_sample_size) ) { 
    Y_top++; X_top++; n++;
  }

  if ( n > max_sample_size ) {
    cerr << "\nError: The maximum sample size = " 
         << max_sample_size 
         << " is too small.\n"
         << "       Give a larger value with the -m flag.\n\n";
    usage_msg();
    exit(1);
  }

  Xsrt = new REAL[n];
  Ysrt = new REAL[n];
  Xsrt_top = Xsrt + n;
  Ysrt_top = Ysrt + n;

  INTEGER * Isrt = new INTEGER[n];
  heapsort(X,n,Xsrt,Isrt);
  for (i=0; i<n; i++) Ysrt[i]=Y[Isrt[i]];

  delete [] Isrt; 

/*
  for (i=0; i<n; i++) 
    cout <<i<<' '<<Ysrt[i]<<' '<<Xsrt[i] <<'\n';
  cout << '\n';
*/

  lo = n/4;  hi = (3*n)/4;
  xspread = *(Xsrt+hi) - *(Xsrt+lo); 
  if (xspread<=0.0) xspread = *(Xsrt+n-1) - *Xsrt;
  xspread /= 1.34;

  if (compute_xmin) xmin=*Xsrt;
  if (compute_xmax) xmax=*(Xsrt+n-1);

  for (x=X, xmean=0.0; x < X_top; x++) {
    xmean += *x; 
  } 
  xmean /= (REAL)n; 
 
  for (x=X, xsigma=0.0; x < X_top; x++) { 
    xsigma += (*x-xmean)*(*x-xmean); 
  } 
  xsigma = sqrt(xsigma/(REAL)(n-1)); 

  delete [] X; 
  delete [] Y; 
}

REAL crossval(int isw)
{
  REAL     *y,    *x,   *w;
  REAL     *yt,   *xt,  *wt, *vt;
  REAL     yplt;
  REAL     res;
  REAL     top, bot, weight, u;
  REAL     cv=0;

  cv=zero;

  if (normal_kernel) {
    for (xt=Xbin,yt=Ybin,wt=Xwgt,vt=Yvar; xt < Xbin_top; yt++,xt++,wt++,vt++) {

      for(top=zero,bot=zero,x=Xbin,y=Ybin,w=Xwgt; x < Xbin_top; x++,y++,w++) {
        if (xt != x) {
          u = (*xt-*x)/bandwidth;
          u *= u;
          weight = (*w) * (REAL) exp((double)(-u/2.0));
          top += *y * weight;
          bot += weight;
        }
      }

      if (bot != zero) {
        yplt = top/bot;
      }
      else {
        yplt = zero;
      }

      res = (*yt - yplt);

      cv += (*wt)*res*res + *vt;

    }
  }

  if (quadratic_kernel) {
  // Fix this. It's inefficient.  There's no need to loop over 
  // the entire sample.

    for (xt=Xbin,yt=Ybin,wt=Xwgt,vt=Yvar; xt < Xbin_top; yt++,xt++,wt++,vt++) {

      for(top=zero,bot=zero,x=Xbin,y=Ybin,w=Xwgt; x < Xbin_top; x++,y++,w++) {
        if (xt != x) {
          u = (*xt-*x)/bandwidth;
          u *= u;
          if ( (double) u <= 1.0 ) {
            weight = (*w) * 0.75 * (1.0 - u);
          }
          else {
            weight = 0.;
          }
          top += *y * weight;
          bot += weight;
        }
      }

      if (bot != zero) {
        yplt = top/bot;
      }
      else {
        yplt = zero;
      }

      res = (*yt - yplt);

      cv += (*wt)*res*res + *vt;
    }
  }

  cv /= (REAL)n;

  if (isw) {
    cout << "Cross validation = " << cv 
         << " for bandwidth = " << bandwidth << "\n";
  }

  return cv;
}

void estimate()
{
  REAL     *y, *x, *w;
  REAL     yplt, xplt;
  REAL     top, bot, weight, u;
  REAL     xinc;

  xinc = (xmax - xmin) / ( (REAL) grid );


  if (normal_kernel) {

    char str[80];

    for (xplt=xmin; xplt < xmax; xplt+=xinc) {

      for(top=zero,bot=zero,x=Xbin,y=Ybin,w=Xwgt; x < Xbin_top; x++,y++,w++) {
        u = (xplt-*x)/bandwidth;
        u *= u;
        weight = (*w) * (REAL) exp((double)(-u/2.0));
        top += *y * weight;
        bot += weight;
      }

      if (bot != zero) {
        yplt = top/bot;
        if (eformat) {
          sprintf(str,"%30.18e %30.18e\n",yplt,xplt); cout << str;
        }
        else {
          cout << yplt << ' ' << xplt << '\n';
        }
      }
      else {
        yplt = zero;
        cerr << "Bandwidth too narrow to estimate at x = " << xplt << '\n';
      }
    }
  }
  
  if (quadratic_kernel) {
  // Fix this. It's inefficient.  There's no need to loop over 
  // the entire sample.

    char str[80];

    for (xplt=xmin; xplt < xmax; xplt+=xinc) {

      for(top=zero,bot=zero,x=Xbin,y=Ybin,w=Xwgt; x < Xbin_top; x++,y++,w++) {
        u = (xplt-*x)/bandwidth;
        u *= u;
        if ( (double) u <= 1.0 ) {
          weight = (*w) * 0.75 * (1.0 - u);
        }
        else {
          weight = 0.;
        }
        top += *y * weight;
        bot += weight;
      }

      if (bot != zero) {
        yplt = top/bot;
        if (eformat) {
          sprintf(str,"%30.18e %30.18e\n",yplt,xplt); cout << str;
        }
        else {
          cout << yplt << ' ' << xplt << '\n';
        }
      }
      else {
        yplt = zero;
        cerr << "Bandwidth too narrow to estimate at x = " << xplt << '\n';
      }
    }
  }
}

void pred()
{
  REAL     *y,    *x,  *w;
  REAL            *xt;
  REAL     yplt;
  REAL     top, bot, weight, u;

  if (normal_kernel) {

    char str[80];

    for (xt=Xsrt; xt < Xsrt_top; xt++) {

      for(top=0,bot=zero,x=Xbin,y=Ybin,w=Xwgt; x < Xbin_top; x++,y++,w++) {
        u = (*xt-*x)/bandwidth;
        u *= u;
        weight = (*w) * (REAL) exp((double)(-u/2.0));
        top += *y * weight;
        bot += weight;
      }

      if (bot != zero) {
        yplt = top/bot;
      }
      else {
        yplt = zero;   //Should never happen.
      }
      if (eformat) {
        sprintf(str,"%30.18e %30.18e\n",yplt,*xt); cout << str;
      }
      else {
        cout << yplt << ' ' << *xt << '\n';
      }
    }
  }

  if (quadratic_kernel) {
  // Fix this. It's inefficient.  There's no need to loop over 
  // the entire sample.

    char str[80];

    for (xt=Xsrt; xt < Xsrt_top; xt++) {

      for(top=0,bot=zero,x=Xbin,y=Ybin,w=Xwgt; x < Xbin_top; x++,y++,w++) {
        u = (*xt-*x)/bandwidth;
        u *= u;
        if ( (double) u <= 1.0 ) {
          weight = (*w) * 0.75 * (1.0 - u);
        }
        else {
          weight = 0.;
        }
        top += *y * weight;
        bot += weight;
      }
    
      if (bot != zero) {
        yplt = top/bot;
      }
      else {
        yplt = zero;  //Should never happen.
      }
      if (eformat) {
        sprintf(str,"%30.18e %30.18e\n",yplt,*xt); cout << str;
      }
      else {
        cout << yplt << ' ' << *xt << '\n';
      }
    }
  }
}

void bin(REAL b) 
{
  const REAL zero = 0.0;
  const REAL one = 1.0;

  INTEGER i,j;

  if ((n <= 50)||(do_not_bin)) {
    Xbin = new REAL[n];
    Ybin = new REAL[n];
    Xwgt = new REAL[n];
    Yvar = new REAL[n];
    Xbin_top = Xbin + n;
    Ybin_top = Ybin + n;
    for (i=0; i<n; i++) {
      Xbin[i]=Xsrt[i];
      Ybin[i]=Ysrt[i];
      Xwgt[i]=one;
      Yvar[i]=zero;
    }
    nbin = n;
/*
    for (i=0; i<nbin; i++) 
      cout <<i<<' '<<Ybin[i]<<' '<<Xbin[i]<<' '<<Xwgt[i]<<' '<<Yvar[i]<<'\n';
    cout << '\n';
*/
    return;
  }

  double fn = (double)n;
  double sn = sqrt(fn);
  double qn = 0.1*fn;
  double d;

  d = (sn < qn) ? sn : qn;
  i = (INTEGER) d;

  INTEGER q10 = ( 0 > i ) ? 0 : i;
  INTEGER q90 = ( n-i < n-1 ) ? n-i : n-1;

  // bins:                 
  // left                         0,    ...,  q10
  // middle                 q10 + 1,    ...,  q10 + nbins_middle,
  // right   q10 + nbins_middle + 1,    ...,  q10 + nbins_middle + n - q90

  INTEGER nbins_middle = (INTEGER) ((Xsrt[q90]-Xsrt[q10])/b);
  INTEGER nbins_total = nbins_middle + q10 + 1 + n - q90;
  INTEGER bin_count = 0;

  Xbin = new REAL[nbins_total];
  Ybin = new REAL[nbins_total];
  Xwgt = new REAL[nbins_total];
  Yvar = new REAL[nbins_total];
  
  REAL tmp;
  tmp = nbins_middle;
  REAL xinc = (Xsrt[q90]-Xsrt[q10])/tmp;

  for (i=0; i<=q10; i++) {
    Xbin[i] = Xsrt[i];
    Ybin[i] = Ysrt[i];
    Xwgt[i] = one;
    Yvar[i] = zero;
    bin_count++;
  }

  tmp = Xsrt[q10];

  for (i=q10+1; i<=q10 + nbins_middle; i++) {
    tmp += xinc;
    Xbin[i] = tmp;
    Ybin[i] = zero;
    Xwgt[i] = zero;
    Yvar[i] = zero;
    bin_count++;
  }

  for (i=q90; i<n; i++) {
    Xbin[q10 + nbins_middle + 1 + i - q90] = Xsrt[i];
    Ybin[q10 + nbins_middle + 1 + i - q90] = Ysrt[i];
    Xwgt[q10 + nbins_middle + 1 + i - q90] = one;
    Yvar[q10 + nbins_middle + 1 + i - q90] = zero;
    bin_count++;
  }

  Xwgt[q10 + nbins_middle + 1] = zero;
  Ybin[q10 + nbins_middle + 1] = zero;

  REAL *x = Xsrt;
  REAL *y = Ysrt;
  REAL low, high;
  for (i=q10+1; i<=q10 + nbins_middle + 1; i++) {
    low = Xbin[i-1];
    high = Xbin[i];
    while (*x<=high) {
      if( (low < *x) && (*x <= high) ) {
        Xwgt[i]++; 
        Ybin[i] += *y;
        Yvar[i] += (*y)*(*y);
      }
      x++; y++;
    }
  }

  if (nbins_total != bin_count) {
    cerr << "Error, "<<progname<< ", bin, this should never happen.";
    exit(1);
  }

  tmp = 0.5*xinc;
  for (i=q10+1; i<=q10 + nbins_middle; i++) {
    Xbin[i] -= tmp;
  }

  for (i=0, j=0; i<bin_count; i++) {
    if (Xwgt[i] > 0 ) {
      Xbin[j] = Xbin[i];
      Ybin[j] = Ybin[i];
      Xwgt[j] = Xwgt[i];
      Yvar[j] = Yvar[i];
      j++;
    }
  }

  bin_count = j;

  Xbin_top = Xbin + bin_count;
  Ybin_top = Ybin + bin_count;
  nbin = bin_count;

  REAL *w;  REAL *v;
  for (y=Ybin, w=Xwgt, v=Yvar; y < Ybin_top; y++, w++, v++) {
    *y /= *w; 
    if (*v > zero) *v -= (*w)*(*y)*(*y); 
  }

/*
  for (i=0; i<bin_count; i++) 
    cout <<i<<' '<<Ybin[i]<<' '<<Xbin[i]<<' '<<Xwgt[i]<<' '<<Yvar[i]<<'\n';
  cout << '\n';
*/
}
