// ArchTOIPipe           (C)     CEA/DAPNIA/SPP IN2P3/LAL
//                               Eric Aubourg
//                               Christophe Magneville
//                               Reza Ansari
#include "config.h"

#include "array.h"
#include "simoffset.h"
#include <math.h>
#include "toimanager.h"
#include "pexceptions.h"
#include "ctimer.h"
#include "ntuple.h"
#include "sopemtx.h"

#include "flagtoidef.h"


SimpleOffsetEstimator::SimpleOffsetEstimator(int mwsz, int nptfit, int degpol)
  : poly((degpol > 1)?degpol+1:2)
{
  SetParams(mwsz, nptfit, degpol);
  totnscount = 0;
  totnbblock = 0;
  bmincut = bmaxcut = -9999.;
  bgalcut = false;
  ns_flgcut = ns_bgalcut = 0;
  npb_fitpoly = 0;
  SavePolyNTuple();
}

SimpleOffsetEstimator::~SimpleOffsetEstimator()
{
}

void SimpleOffsetEstimator::SetParams(int mwsz, int nptfit, int degpol)
{
  mWSz = (mwsz > 8) ? mwsz : 8;
  nPtFit = (nptfit > degpol+2) ? nptfit : degpol+2;
  if (nPtFit%2 == 0) nPtFit++;
  degPol = (degpol > 1)?degpol:1;
  
}

void SimpleOffsetEstimator::SetBGalCut(double bmin, double bmax)
{
  bgalcut = true;
  bmincut = bmin;
  bmaxcut = bmax;
}

void SimpleOffsetEstimator::PrintStatus(::ostream & os)
{
  os << "\n ------------------------------------------------------ \n" 
     << " SimpleOffsetEstimator::PrintStatus() - MeanWSize= " << mWSz << " NPtFit=" 
     << nPtFit << " DegPoly=" << degPol << endl;
  if (bgalcut) 
    os << " bGalCut = Yes , bGalMin = " << bmincut << " bGalMax= " << bmaxcut << endl;
  else os << " bGalCut = No " << endl;
  TOIProcessor::PrintStatus(os);
  os << " ProcessedSampleCount=" << ProcessedSampleCount() 
     << " NS_FlgCut= " << ns_flgcut << " NS_BGalCut= " << ns_bgalcut 
     << " NPb_FitPoly= " << npb_fitpoly << endl;
  os << " ------------------------------------------------------ " << endl;  
}

void SimpleOffsetEstimator::init()
{
  cout << "SimpleOffsetEstimator::init" << endl;
  declareInput("in");
  declareInput("bgal");
  declareOutput("offset");
  declareOutput("out");
  declareOutput("incopie");
  declareOutput("bgalcopie");
  //  declareOutput("mean_y");
  //  declareOutput("sig_y");
  //  declareOutput("mean_x");
  name = "SimpleOffsetEstimator";
}

void SimpleOffsetEstimator::run()
{
  int snb = getMinIn();
  int sne = getMaxIn();

  bool fgoffset = checkOutputTOIIndex(0);
  bool fgbgal = checkInputTOIIndex(1);
  bool fgout = checkOutputTOIIndex(1);
  bool fgincopie = checkOutputTOIIndex(2);
  bool fgbgalcopie = checkOutputTOIIndex(3);
  //  bool fgmeany = checkOutputTOIIndex(4);
  //  bool fgsigy = checkOutputTOIIndex(5);
  //  bool fgmeanx = checkOutputTOIIndex(6);

  if (!checkInputTOIIndex(0)) {
    cerr << " SimpleOffsetEstimator::run() - Input TOI (in) not connected! "
	 << endl;
    throw ParmError("SimpleOffsetEstimator::run() Input TOI (in) not connected!");
  }
  if (bgalcut && !fgbgal) {
    cerr << " SimpleOffsetEstimator::run() - Input TOI bgal not connected! "
	 << endl;
    throw ParmError("SimpleOffsetEstimator::run() Input TOI bgal not connected!");
  }
  if (fgbgalcopie && !fgbgal) {
    cerr << " SimpleOffsetEstimator::run() - Output TOI bgalcopie connected without input TOI bgal!"
	 << endl;
    throw ParmError("SimpleOffsetEstimator::run() Output TOI bgalcopie connected without input TOI bgal!");
  }
  if (!fgoffset && !fgout) {
    cerr << " SimpleOffsetEstimator::run() - No output TOI (offset/in-offset) connected!"
	 << endl;
    throw ParmError(" SimpleOffsetEstimator::run() No output TOI (offset/in-offset) connected!");
  }
  
  cout << " SimpleOffsetEstimator::run() SNRange=" << snb << " - " << sne << endl; 

  // NTuple pour sauvegarde des coeff de poly  
  char * nomsnt[] = {"sncur", "sn0", "meanx", "meany", "sigy", 
		       "a0", "a1", "a2", "ycur", "nok", "nbblkok"};
  NTuple xntp(11, nomsnt);

  char ans[20];

  try {

    // Vecteurs pour les donnees et les sorties
    int wsize = mWSz;

    int hisblk = nPtFit/2+1;
    Vector vinhist(wsize*hisblk);
    TVector<uint_8> vfghist(wsize*hisblk);
    Vector voff(wsize);
    Vector vout(wsize);
    Vector vinc(wsize);
    TVector<uint_8> vfgc(wsize);
    Vector bgal(wsize);

    
    // Pour le fit 
    Vector errCoef(degPol+1);
    
    Vector X(nPtFit+1);
    Vector X0(nPtFit+1);
    Vector Y(nPtFit+1);
    Vector YErr(nPtFit+1);

    // Variables diverses 
    int k,kb,j,klast;
    klast = snb-1;
    totnbblock = 0;
    

    int nbblkok = 0;

    bool fginiXYdone = false;

    double sn0 = 0.;
    double nok = 0.;
    double mean = 0.;
    double sig = 0.;
    double meanx = 0.;
    double sn_last = 0.;
    double y_last = 0.;

    double last_meanok = 0.;
    double last_sigok = 1.;
    double meanx_forlast = 0.;
    bool fg_last_meansig = false;
 
    // Boucle sur les sampleNum
    // 1ere partie, on traite par paquets de wsize
    int nblk = (sne-snb+1)/wsize; 
    for(kb=0; kb<nblk; kb++) {
      //      cout << " SimpleOffsetEstimator::run - Loop " << kb << " NbBlkOK= "  
      //	   << nbblkok << " <CR> to continue / q --> QUIT ... \n" ; 
      //     gets(ans); if (ans[0] == 'q') break;
      k = kb*wsize+snb;
      //    for(k=snb;k<=sne-wsize+1;k+=wsize) {
      // Lecture d'un bloc de donnees
      Vector vin( vinhist(Range((kb%hisblk)*wsize, 0, wsize) ) );
      TVector<uint_8> vfg( vfghist(Range((kb%hisblk)*wsize, 0, wsize) ) );

      getData(0, k, wsize, vin.Data(), vfg.Data());
      if (fgbgal) {
	getData(1, k, wsize, bgal.Data());
	if (fgbgalcopie) putData(3, k, wsize, bgal.Data());
      }

      fg_last_meansig = false;

      if (kb == 0) {
	last_meanok = vin.Sum()/wsize;
	last_sigok = vin.SumX2()/wsize - last_meanok*last_meanok;
      }

      // Calcul moyenne et sigma du bloc 
      nok = 0.;  meanx = 0.;
      mean = 0.;  sig = 0.;
      meanx_forlast = 0.;
      for(j=0; j<wsize; j++) {
	meanx_forlast += k+j;
	if ( vfg(j) ) { ns_flgcut++;  continue; }
	if (bgalcut && (bgal(j) > bmincut) && (bgal(j) < bmaxcut) ) 
	  { ns_bgalcut++;  continue; }
	mean += vin(j);
	sig += vin(j)*vin(j);
	meanx += k+j;
	nok++;
      }
      
      
      if (!fginiXYdone) {
	if (nok > 0.5) {
	  mean /= nok;
	  meanx /= nok;
	  sig = sig/nok-mean*mean;	
	}
	X = RegularSequence((kb+0.5)*wsize, (double)wsize);
	Y = mean;
	YErr = (nok > 0.5) ? sqrt(mean) : 1.;
	fginiXYdone = true; 
      }

      if (nok > 3.5) {
	mean /= nok;
	meanx /= nok;
	sig = sig/nok-mean*mean;
	if (sig < 1.e-10*mean) sig = 1.e-10*mean;
	if (sig < 1.e-39) sig = 1.e-39;
	int kk = nbblkok%nPtFit;
	nbblkok++;
	Y(kk) = mean; 
	YErr(kk) = sig;
	X(kk) = meanx;
	last_meanok = mean;
	last_sigok = sig;
      }
      else {
	int kk = nbblkok%nPtFit;
	Y(kk) = last_meanok; 
	YErr(kk) = last_sigok*10.;
	X(kk) = meanx_forlast/wsize;	
	fg_last_meansig = true;
      }

            
      if (((nok>3.5) || fg_last_meansig) && (nbblkok >= nPtFit) ) {
	// On force le fit a garder une certaine continuite
	Y(nPtFit) = y_last;
	double smin, smax;
	YErr(Range(0,0,nPtFit)).MinMax(smin, smax);
	if (smax < 1.e-39)   smax = 1.e-39;
	YErr(nPtFit) = smax/10.;
	X(nPtFit) = sn_last;
	sn0 = (double)(k-nPtFit*wsize*0.5);
	X0 = X;
	X0 -= sn0;
	FitPoly(X0, Y, YErr);
      }
      else {
	if (nbblkok < 2) {
	  sn0 = k+1;
	  poly(0) = mean;
	  for(int jj=1; jj<=degPol; jj++) poly(jj) = 0.;
	}
	else if (nbblkok < nPtFit) {
	  poly(0) = 0.5*(Y(nbblkok-1)+Y(0));
	  poly(1) = (Y(nbblkok-1) - Y(0))/(X(nbblkok-1)-X(0));
	  sn0 =  0.5*(X(nbblkok-1)+X(0));
	  for(int jj=2; jj<=degPol; jj++) poly(jj) = 0.;
	} 
	sn_last = sn0;
	y_last = poly(sn_last-sn0);
      }
      
      if (ntpoly) {    // Remplissage du XNTuple de controle 
	float xnt[12];
	xnt[0] = k;
	xnt[1] = sn0;
	xnt[2] = meanx;
	xnt[3] = mean;
	xnt[4] = sig;
	xnt[5] = poly(0);
	xnt[6] = poly(1);
	xnt[7] = poly(2);
	xnt[8] = ApplyPoly(k-sn0);
	xnt[9] = nok;
	xnt[10] = nbblkok;
	xntp.Fill(xnt);
      }

     
      if (fgincopie) putData(2, k, wsize, vin.Data(), vfg.Data());

      if (kb < nPtFit/2) continue;
      Vector vinc; 
      TVector<uint_8> vfgc;
      int kbs = kb-nPtFit/2;
      k = kbs*wsize+snb;
      if (kb == nblk-1) {
	int wszt = wsize*hisblk;
	voff.ReSize(wszt);
	vout.ReSize(wszt);	
	vinc.ReSize(wszt); 
	vfgc.ReSize(wszt);
	bgal.ReSize(wszt); 
	int jb = 0;
	for(int kbsc=kbs; kbsc<nblk; kbsc++) {
	  vinc(Range(jb*wsize, 0, wsize)) = 
	       vinhist(Range((kbsc%hisblk)*wsize, 0, wsize) ) ;
	  vfgc(Range(jb*wsize, 0, wsize)) = 
	       vfghist(Range((kbsc%hisblk)*wsize, 0, wsize) ) ;
	  jb++;     
	}
	wsize = wszt;
      }
      else {
	vinc = vinhist(Range((kbs%hisblk)*wsize, 0, wsize) ) ;
	vfgc = vfghist(Range((kbs%hisblk)*wsize, 0, wsize) ) ;
      }

      // Calcul des valeurs d'offset en sortie
      for(j=0; j<wsize; j++) 
	voff(j) = ApplyPoly(k+j-sn0);
      sn_last = k+wsize;
      y_last = ApplyPoly(sn_last-sn0);
      if (fgoffset) putData(0, k, wsize, voff.Data());
      if (fgout) { 
	vinc -= voff;
	putData(1, k, wsize, vinc.Data(), vfgc.Data());
      }

      klast+=wsize;
      totnscount+=wsize;
      totnbblock++;
      //      cout << " SimpleOffset::run " << totnbblock << " totnscount " << totnscount << endl;
    } // Fin boucle sur les samples, par pas de wsize

    // 3eme partie, on traite la fin du bloc d'echantillons si necessaire
    if (klast < sne) {
      wsize = sne-klast;
      Vector vin(wsize);
      voff.ReSize(wsize);
      vout.ReSize(wsize);
      TVector<uint_8> vfg(wsize);
      k = klast+1;
      getData(0, k, wsize, vin.Data(), vfg.Data());

      if (fgbgal) {
	getData(1, k, wsize, bgal.Data());
	if (fgbgalcopie) putData(3, k, wsize, bgal.Data());
      }

      for(j=0; j<wsize; j++) 
	voff(j) = ApplyPoly(k+j-sn0);
      if (fgoffset) putData(0, k, wsize, voff.Data());
      if (fgincopie) putData(2, k, wsize, vin.Data(), vfg.Data());
      if (fgout) { 
	vin -= voff;
	putData(1, k, wsize, vin.Data(), vfg.Data());
      }

      /*      
      if (fgmeany) {
	vout = mean;
	putData(4, k, wsize, vout.Data());
      }
      if (fgsigy) {
	vout = sig;
	putData(5, k, wsize, vout.Data());
      }
      
      if (fgmeanx) {
	vout = meanx;
	putData(6, k, wsize, vout.Data());
      }
      */
      
      klast+=wsize;
      totnscount+=wsize;
      totnbblock++;      
    }
    
    cout << " SimpleOffsetEstimator::run() - End of processing " 
	 << " TotNbBlocks= " << totnbblock << " ProcSamples=" << totnscount << endl;
    
  }  // Bloc try 

  catch (PException & exc) {
    cerr << "SimpleOffsetEstimator::run() Catched Exception " << (string)typeid(exc).name()
         << "\n .... Msg= " << exc.Msg() << endl;
  }

  if (ntpoly) {
    if (ntpolyname.length() < 1)  ntpolyname = "simoffset.ppf";
    POutPersist pos(ntpolyname);
    cout << " SimpleOffsetEstimator::run()/Info : Writing poly ntuple to PPF file " << ntpolyname << endl;
    pos << xntp;
  }
}


double SimpleOffsetEstimator::ApplyPoly(double x)
{
  if (degPol == 0) return(poly(0));
  else if (degPol == 1)  return(poly(0)+poly(1)*x);
  else if (degPol == 2)  return(poly(0)+poly(1)*x+poly(2)*x*x);
  else if (degPol == 3)  return(poly(0)+poly(1)*x+poly(2)*x*x+poly(3)*x*x*x);
  else {
    double ret = poly(0);
    double xx = x;
    for(int k=1; k<degPol+1; k++)  { ret += xx*poly(k);  xx *= x; }
    return(ret);
  }
}

// Fonction pour faire LinFit d'un polynome
r_8 rzpoly_f(uint_4 k, r_8 x )
{
  if (k == 0) return(1.);
  else if (k == 1) return(x);
  else if (k == 2) return(x*x);
  else if (k == 3) return(x*x*x);
  else {
    r_8 ret = x;
    for(int i=1; i<k; i++) ret *= x;
    return(ret);
  }
}

void SimpleOffsetEstimator::FitPoly(Vector& X0, Vector& Y, Vector& YErr)
{
  Vector cpol(degPol+1);
  Vector errcoef(degPol+1);
  LinFitter<r_8> lfit;
  try {
    lfit.LinFit(X0, Y, YErr, degPol+1, rzpoly_f, cpol, errcoef);
    poly = cpol;
  }
  catch (PException & exc) {
    if (npb_fitpoly < 50) 
      cout << " -- SimpleOffsetEstimator::FitPoly/ Catched Exception " 
	   << (string)typeid(exc).name() << "\n .... Msg= " << exc.Msg() << endl;
    if (npb_fitpoly < 10) 
      cout << " X0= " << X0 << " Y=" << Y << " YErr=" << YErr << endl;
    npb_fitpoly++;
  }
  return;
}