// polfitclip.cc
// Eric Aubourg         CEA/DAPNIA/SPP   novembre 1999

#include <iostream.h>

#ifndef ANSI
#define ANSI
#endif

#include <math.h>
#include <string>
#include <algorithm>

extern "C" {
#include "nrutil.h"
}

#include "polfitclip.h"

PolFitClip::PolFitClip(long ndatamax, long degpol) {
  this->ndata = 0;
  this->ndatamax = ndatamax;
  this->nparm   = degpol+1;
  
  x    = ::dvector(1, ndatamax);
  y    = ::dvector(1, ndatamax);
  ySort= ::dvector(1, ndatamax);
  sig  = ::dvector(1, ndatamax);
  a    = ::dvector(1, nparm);
  ia   = ::ivector(1, nparm);
  cov  = ::matrix(1, nparm, 1, nparm);
 
  for (int i=1; i<=ndatamax; i++) sig[i] = 1;

  for (int i=1; i<=nparm; i++) ia[i] = 1;

  frc   = 0;
  nsig  = 3;
  npass = 3;
}

PolFitClip::~PolFitClip() {
  free_vector(x,     1, ndatamax);
  free_vector(y,     1, ndatamax);
  free_vector(ySort, 1, ndatamax);
  free_vector(sig,   1, ndatamax);
  free_vector(a,     1, nparm);
  free_ivector(ia,   1, ndatamax);
  free_matrix(cov,   1, nparm, 1, nparm);
}

void PolFitClip::setData(double const* xx, double const* yy, long n) {
  if (n>ndatamax) {
    cerr << "PolFitClip::setData : n > ndatamax " << n << " " << ndatamax << endl;
    exit(-1);
  }
  ndata = n;
  for (int i=1; i<=n; i++) {
    x[i] = xx[i-1];
    y[i] = yy[i-1];
  }
}

void PolFitClip::addData(double xx, double yy) {
  ndata++;
  if (ndata>ndatamax) {
    cerr << "PolFitClip::addData : n > ndatamax " << ndata << " " << ndatamax << endl;
    exit(-1);
  }
  x[ndata] = xx;
  y[ndata] = yy;
}

void PolFitClip::clear() {
  this->ndata = 0;
}

void PolFitClip::setClip(double frc, double nsig, int npass) {
  this->frc   = frc;
  this->nsig  = nsig;
  this->npass = npass;
}

extern "C" void PolFitClip_polfunc(double x, double afunc[], int ma);

void PolFitClip_polfunc(double x, double afunc[], int ma) {
  afunc[1] = 1;
  for (int i=2; i<=ma; i++)
    afunc[i] = afunc[i-1]*x;
}

extern "C"
void lfit(double x[], double y[], double sig[], int ndat, double a[], int ia[],
	int ma, double **covar, double *chisq, void (*funcs)(double, double [], int));

int PolFitClip::doFit1() {
  if (ndata < nparm) return -1;

  try {
    lfit(x, y,  sig, ndata, a, ia, nparm, cov, &chi2, PolFitClip_polfunc);
  } catch (string st) {
    return -2;
  }
  
  return 0;
}

int PolFitClip::doFit(double* coef) {
  if (ndata < nparm) return -1;
  
  memcpy(ySort, y, (ndata+1)*sizeof(double));
  sort(ySort+1, ySort+ndata+1);
  int i = frc*ndata;
  double cutMin = ySort[i+1];
  i = ndata-i;
  double cutMax = ySort[i];
  
  sigma = 0;
  
  for (int ipass = 0; ipass<npass; ipass++) {
    for (int i=1; i<=ndata; i++) {
      if (y[i] >= cutMin && y[i] <= cutMax)
        sig[i]=1;
      else
        sig[i]=9.e99;
    }
    double s=0, s2=0;
    int n=0;
    for (int i=1; i<=ndata; i++) {
      if (sig[i]>=10) continue;
      double yy = ipass ? value(x[i]) : 0;
      if (ipass>0) {
        if (fabs(yy-y[i]) > sigma*nsig) {
          sig[i]=9.e99;
          continue;
        }
      }
      s += yy-y[i];
      s2 +=  (yy-y[i])*(yy-y[i]);
      n++;
    }
    if (n<1) return -3;
    s/=n;
    s2/=n;
    sigma = s2 - s*s;
    sigma = (sigma>0) ? sqrt(sigma) : 0;
    
    int rc = doFit1(); 
    if (rc) return rc;
  }
  
  if (coef)
    for (int i=1; i<=nparm; i++) {
      coef[i-1] = a[i];
    }
  
  return 0;
}


double PolFitClip::value(double x) {
  double r = a[nparm];
  for (int i=nparm-1; i>0; i--) {
    r = r*x+a[i];
  }
  return r;
}


long PolFitClip::getNData() {
  return ndata;
}

double PolFitClip::getXMin() {
  double xm = 9.e99;
  for (int i=1; i<=ndata; i++)
    if (x[i]<xm) xm=x[i];
  return xm;
}

double PolFitClip::getXMax() {
  double xm = -9.e99;
  for (int i=1; i<=ndata; i++)
    if (x[i]>xm) xm=x[i];
  return xm;
}
  
double PolFitClip::getChi2() {
  return chi2;
}

double PolFitClip::getSigma() {
  return sigma;
}

double PolFitClip::getCov(int i, int j) {
  return cov[i+1][j+1];
}


/************************************************/

PolFitClip2::PolFitClip2(long ndatamax, long degpol) {
  this->ndata = 0;
  this->ndatamax = ndatamax;
  this->nparm   = degpol+1;
  
  x    = ::dvector(1, ndatamax);
  y    = ::dvector(1, ndatamax);
  ySort= ::dvector(1, ndatamax);
  z    = ::dvector(1, ndatamax);
  zSort= ::dvector(1, ndatamax);
  sig  = ::dvector(1, ndatamax);
  aY   = ::dvector(1, nparm);
  aZ   = ::dvector(1, nparm);
  ia   = ::ivector(1, nparm);
  covY = ::matrix(1, nparm, 1, nparm);
  covZ = ::matrix(1, nparm, 1, nparm);
 
  for (int i=1; i<=ndatamax; i++) sig[i] = 1;

  for (int i=1; i<=nparm; i++) ia[i] = 1;

  frcY  = 0;
  frcZ  = 0;
  nsig  = 3;
  npass = 3;
}

PolFitClip2::~PolFitClip2() {
  free_vector(x,     1, ndatamax);
  free_vector(y,     1, ndatamax);
  free_vector(ySort, 1, ndatamax);
  free_vector(z,     1, ndatamax);
  free_vector(zSort, 1, ndatamax);
  free_vector(sig,   1, ndatamax);
  free_vector(aY,    1, nparm);
  free_vector(aZ,    1, nparm);
  free_ivector(ia,   1, ndatamax);
  free_matrix(covY,  1, nparm, 1, nparm);
  free_matrix(covZ,  1, nparm, 1, nparm);
}

void PolFitClip2::setData(double const* xx, double const* yy, double const* zz, long n) {
  if (n>ndatamax) {
    cerr << "PolFitClip2::setData : n > ndatamax " << n << " " << ndatamax << endl;
    exit(-1);
  }
  ndata = n;
  for (int i=1; i<=n; i++) {
    x[i] = xx[i-1];
    y[i] = yy[i-1];
    z[i] = zz[i-1];
  }
}

void PolFitClip2::addData(double xx, double yy, double zz) {
  ndata++;
  if (ndata>ndatamax) {
    cerr << "PolFitClip2::addData : n > ndatamax " << ndata << " " << ndatamax << endl;
    exit(-1);
  }
  x[ndata] = xx;
  y[ndata] = yy;
  z[ndata] = zz;
}

void PolFitClip2::clear() {
  this->ndata = 0;
}

void PolFitClip2::setClip(double frcY, double frcZ, double nsig, int npass) {
  this->frcY  = frcY;
  this->frcZ  = frcZ;
  this->nsig  = nsig;
  this->npass = npass;
}


int PolFitClip2::doFit1() {
  if (ndata < nparm) return -1;

  try {
    lfit(x, y,  sig, ndata, aY, ia, nparm, covY, &chi2Y, PolFitClip_polfunc);
    lfit(x, z,  sig, ndata, aZ, ia, nparm, covZ, &chi2Z, PolFitClip_polfunc);
  } catch (string st) {
    return -2;
  }
  
  return 0;
}

int PolFitClip2::doFit(double* coefY, double* coefZ) {
  if (ndata < nparm) return -1;
  
  memcpy(ySort, y, (ndata+1)*sizeof(double));
  sort(ySort+1, ySort+ndata+1);
  memcpy(zSort, z, (ndata+1)*sizeof(double));
  sort(zSort+1, zSort+ndata+1);
  
  int i = frcY*ndata;
  double cutMinY = ySort[i+1];
  i = ndata-i;
  double cutMaxY = ySort[i];
  
  i = frcZ*ndata;
  double cutMinZ = zSort[i+1];
  i = ndata-i;
  double cutMaxZ = zSort[i];
  
  sigmaY = 0;
  sigmaZ = 0;
  
  for (int ipass = 0; ipass<=npass; ipass++) {
    for (int i=1; i<=ndata; i++) {
      if (y[i] >= cutMinY && y[i] <= cutMaxY &&
          z[i] >= cutMinZ && z[i] <= cutMaxZ)
        sig[i]=1;
      else
        sig[i]=9.e99;
    }
    double sy=0, sy2=0, sz=0, sz2=0;
    ndataused=0;
    for (int i=1; i<=ndata; i++) {
      if (sig[i]>=10) continue;
      double yy = ipass ? valueY(x[i]) : 0;
      double zz = ipass ? valueZ(x[i]) : 0;
      if (ipass>0 && ipass < npass) {
        if (fabs(yy-y[i]) > sigmaY*nsig ||
            fabs(zz-z[i]) > sigmaZ*nsig) {
          sig[i]=9.e99;
          continue;
        }
      }
      sy += yy-y[i];
      sy2 +=  (yy-y[i])*(yy-y[i]);
      sz += zz-z[i];
      sz2 +=  (zz-z[i])*(zz-z[i]);
      ndataused++;
    }
    if (ndataused<nparm) return -3;
    sy/=ndataused;
    sy2/=ndataused;
    sigmaY = sy2 - sy*sy;
    sigmaY = (sigmaY>0) ? sqrt(sigmaY) : 0;
    sz/=ndataused;
    sz2/=ndataused;
    sigmaZ = sz2 - sz*sz;
    sigmaZ = (sigmaZ>0) ? sqrt(sigmaZ) : 0;
    if (ipass == npass) break;
    
    int rc = doFit1(); 
    if (rc) return rc;
  }
  
  if (coefY)
    for (int i=1; i<=nparm; i++) {
      coefY[i-1] = aY[i];
    }
  if (coefZ)
    for (int i=1; i<=nparm; i++) {
      coefZ[i-1] = aZ[i];
    }
  
  return 0;
}


double PolFitClip2::valueY(double x) {
  double r = aY[nparm];
  for (int i=nparm-1; i>0; i--) {
    r = r*x+aY[i];
  }
  return r;
}

double PolFitClip2::valueZ(double x) {
  double r = aZ[nparm];
  for (int i=nparm-1; i>0; i--) {
    r = r*x+aZ[i];
  }
  return r;
}


long PolFitClip2::getNData() {
  return ndata;
}

long PolFitClip2::getNDataUsed() {
  return ndataused;
}

double PolFitClip2::getXMin() {
  double xm = 9.e99;
  for (int i=1; i<=ndata; i++)
    if (x[i]<xm) xm=x[i];
  return xm;
}

double PolFitClip2::getXMax() {
  double xm = -9.e99;
  for (int i=1; i<=ndata; i++)
    if (x[i]>xm) xm=x[i];
  return xm;
}
  
double PolFitClip2::getChi2Y() {
  return chi2Y;
}

double PolFitClip2::getChi2Z() {
  return chi2Z;
}

double PolFitClip2::getSigmaY() {
  return sigmaY;
}

double PolFitClip2::getSigmaZ() {
  return sigmaZ;
}

double PolFitClip2::getCovY(int i, int j) {
  return covY[i+1][j+1];
}

double PolFitClip2::getCovZ(int i, int j) {
  return covZ[i+1][j+1];
}