#include "segyops.hh"


// mutefun = 0, t \le 0; mutefun = 1, t \ge 1
float mutefun(float t) {
    t = MIN(1.0f,MAX(t,0.0f));
    return 3*t*t-2*t*t*t;
}

// returns true if a < x < b, sets fac=hermite cubic,
// else returns false and fac=1
int taperfun(int lr, float a, float b, float x, float & fac) {
  if (x<=a) {
    if (lr==0) {
      fac*=0.0f;
      return 0;
    }
    else {
      fac*=1.0f;
      return 1;
    }
  }
  else if (x>=b) {
    if (lr==0) {
      fac*=1.0f;
      return 1;
    }
    else {
      fac*=0.0f;
      return 0;
    }
  }
  else {
    if (lr==0) fac*=mutefun((x-a)/(b-a));
    else fac*=mutefun((b-x)/(b-a));
    return 0;
  }
}

float cosfun(float t){
    return 0.5 - cos(M_PI * t)/2.0f;
}
float cosfun2(float t){
    return 0.5 + cos(M_PI * t)/2.0f;
}


// precondition for all of these FOs: input, output have same geom

namespace TSOpt {
    
  using RVL::LocalDataContainer;
  using RVL::RVLException;


void conv(int shift, int nout, int nin, int nker,
	  float * restrict out,
	  const float * restrict in,
	  const float * restrict ker,
	  float scal ) {
  int start;
  int stop;
  for (int i=0; i<nout; i++) {
    start = iwave_max(i+shift-nker+1,0);
    stop  = iwave_min(i+shift+1,nin);
    out[i]=0.0f;

    //#pragma ivdep
    for (int j=start; j<stop; j++) {
      out[i] += ker[i-j+shift]*in[j];
    }
    out[i] *= scal;
  }
}

void corr(int shift, int nout, int nin, int nker,
	  float * restrict out,
	  const float * restrict in,
	  const float * restrict ker,
	  float scal ) {
  int start;
  int stop;
  for (int i=0; i<nout; i++) {
    start = iwave_max(i-shift,0);
    stop  = iwave_min(i-shift+nker,nin);
    out[i]=0.0f;
    //#pragma ivdep
    for (int j=start; j<stop; j++) {
      out[i] += ker[j-i+shift]*in[j];
    }
    out[i] *= scal;
  }
}


  /* in physical discrete coordinates with sample rate dt, convolution
   computes out(ip), outbeg <= ip <= outend, as

   sum from jp = max(ip-kerend, inbeg) to min(ip-kerbeg, inend) 
   of ker(ip-jp)in(jp) = ker(kp)in(jp), 

   array indices: ip=i+outbeg, jp=j+inbeg, kp=k+kerbeg=i-j+outbeg-inbeg
   so kernel index = k = i-j+outbeg-inbeg-kerbeg
   
   so 
   max(i+outbeg-kerend,inbeg) <= j+inbeg <= min(i+outbeg-kerbeg,inend)
   max(i+oubbeg-kerend,inbeg)-inbeg <= j <= min(i+outbeg-kerbeg,inend)-inbeg 

   0 <= i <= outend-outbeg

   check: k=i-j+outbeg-inbeg-kerbeg >= i - min(i+outbeg-kerbeg,inend)+inbeg+outbeg-inbeg-kerbeg
   >= i+outbeg-kerbeg - min(i+outbeg-kerbeg,inend) >=0 

   sim k <= kerend-kerbeg

   in terms of start and length
   nout = outend-outbeg+1  : outend = nout+outbeg-1
   nin  = inend-inbeg+1    : inend  = nin +inbeg -1
   nker = kerend-kerbeg+1  : kerend = nker+kerbeg-1

   so: 0 <= i < nout
   max(i+outbeg-nker-kerbeg+1,inbeg)-inbeg <= j < min(i+outbeg-kerbeg,nin+inbeg-1)-inbeg+1
   max(i+outbeg-kerbeg-inbeg+1-nker,0)     <= j < min(i+outbeg-kerbeg-inbeg+1,nin)
*/

  void SEGYConvolve::operator()(LocalDataContainer<float> & x,
				LocalDataContainer<float> const & y,
				LocalDataContainer<float> const & z) {
    try {
      segytrace & cpout = dynamic_cast<segytrace &>(x);
      segy & trout = cpout.getMetadata();
      segytrace const & cpin = dynamic_cast<segytrace const &>(y);
      segy const & trin = cpin.getMetadata();
      segytrace const & cpker = dynamic_cast<segytrace const &>(z);
      segy const & trker = cpker.getMetadata();

      Value val;
      float dt;
      int nin;
      int nout;
      int nker;
      
      int inbeg;
      int outbeg;
      int kerbeg;
      
      string dtstr="dt";
      string nsstr="ns";
      string rtstr="delrt";

      gethdval(&trout,  (char*)(dtstr.c_str()), &val);
      dt=vtof(hdtype((char*)(dtstr.c_str())), val);

      gethdval(&trin,  (char*)(nsstr.c_str()), &val);
      nin=vtoi(hdtype(nsstr.c_str()), val);
      gethdval(&trin,  (char*)(rtstr.c_str()), &val);
      inbeg=int(0.01 + vtof(hdtype(rtstr.c_str()), val)/dt);

      gethdval(&trout,  (char*)(nsstr.c_str()), &val);
      nout=vtoi(hdtype(nsstr.c_str()), val);
      gethdval(&trout,  (char*)(rtstr.c_str()), &val);
      outbeg=int(0.01 + vtof(hdtype(rtstr.c_str()), val)/dt);

      gethdval(&trker,  (char*)(nsstr.c_str()), &val);
      nker=vtoi(hdtype(nsstr.c_str()), val);
      gethdval(&trker,  (char*)(rtstr.c_str()), &val);
      kerbeg=int(0.01 + vtof(hdtype(rtstr.c_str()), val)/dt);

      if (adj==0) {
	int ishift=outbeg-inbeg-kerbeg;
	conv(ishift,nout,nin,nker,
	     trout.data,trin.data,trker.data);
      }
      else {
	int ishift=inbeg-kerbeg-outbeg;
	corr(ishift,nout,nin,nker,
	     trout.data,trin.data,trker.data);	
      }
    }
    catch (bad_cast) {
      RVLException e;
      e<<"ERROR: SEGYConvolve::operator()\n";
      e<<"  at least one input not segytrace type\n";
      throw e;
    }
  }

  void SEGYTaper::operator()(LocalDataContainer<float> & x,
			     LocalDataContainer<float> const & y) {
    try {
      segytrace & cpout = dynamic_cast<segytrace &>(x);
      segy & trout = cpout.getMetadata();
      segytrace const & cpin = dynamic_cast<segytrace const &>(y);
      segy const & trin = cpin.getMetadata();
      Value val;
      float fac=1.0f;
      int dontdoit = 1;
      // fetch coordinates for each taper axis
      float xcoord;
      
      std::string tmp = "scalco";
      gethdval(&trout,(char *)(tmp.c_str()),&val);
      float scalco = vtof(hdtype("scalco"),val);
      for (int i=0; i<tap.getSize(); i++) {
	// special case: offset, computed from source/receiver posns
	if (tap[i].key=="offset") {
	  tmp="sx";
	  gethdval(&trout,(char *)(tmp.c_str()),&val);
	  float sx = vtof(hdtype("sx"),val);	  
	  tmp="sy";
	  gethdval(&trout,(char *)(tmp.c_str()),&val);
	  float sy = vtof(hdtype("sy"),val);	  
	  tmp="gx";
	  gethdval(&trout,(char *)(tmp.c_str()),&val);
	  float gx = vtof(hdtype("gx"),val);	  
	  tmp="gy";
	  gethdval(&trout,(char *)(tmp.c_str()),&val);
	  float gy = vtof(hdtype("gy"),val);	  
	  xcoord = sqrt(((sx-gx)*(sx-gx))+((sy-gy)*(sy-gy)));
	}
	else {
	  gethdval(&trout,(char *)(tap[i].key.c_str()),&val);
	  xcoord = vtof(hdtype(tap[i].key.c_str()),val);
	}
	if (scalco > 0) { xcoord *=  scalco; }
	if (scalco < 0) { xcoord /= -scalco; }
	dontdoit*=taperfun(0,tap[i].t[0],tap[i].t[1],xcoord,fac); // returns 1 if x to r of t[1] 
	dontdoit*=taperfun(1,tap[i].t[2],tap[i].t[3],xcoord,fac); // returns 1 if x to l of t[2]
	//	cerr<<"xcoord="<<xcoord<<" t0="<<tap[i].t[0]<<" t1="<<tap[i].t[1]<<" t2="<<tap[i].t[2]<<" t3="<<tap[i].t[3]<<" dontdoit="<<dontdoit<<" fac="<<fac<<endl;
      }
      tmp="ns";
      gethdval(&trout,(char *)(tmp.c_str()),&val);
      int ns = vtoi(hdtype("ns"),val);
      if (dontdoit) {
	for (int k=0;k<ns;k++) (trout.data)[k] = (trin.data)[k];
      }
      else {
	for (int k=0;k<ns;k++) (trout.data)[k] = fac*((trin.data)[k]);
      }
    }
    catch (bad_cast) {
      RVL::RVLException e;
      e<<"ERROR: SEGYTaper::operator[]\n";
      e<<"  input not of segytrace type\n";
      throw e;
    }
  }
      
  void SEGYLinMute::operator()(LocalDataContainer<float> & x,
                                 LocalDataContainer<float> const & y) {
        
      try {
            
	segytrace & cpout = dynamic_cast<segytrace &>(x);
	segytrace const & cpin = dynamic_cast<segytrace const &>(y);
            
	segy & trout = cpout.getMetadata();
	segy const & trin = cpin.getMetadata();
            
	Value val;
	float gx;
	float x;
	float dt;
	float t0;
	float wm;
	int nt;
            
	string name="ns";
	gethdval(&trin,(char *)(name.c_str()),&val);
	nt=vtoi(hdtype((char *)(name.c_str())),val);
            
	name="dt";
	gethdval(&trin,(char *)(name.c_str()),&val);
	dt=0.001*vtof(hdtype((char *)(name.c_str())),val);
            
	name="delrt";
	gethdval(&trin,(char *)(name.c_str()),&val);
	t0=vtof(hdtype((char *)(name.c_str())),val);
            
	name="gx";
	gethdval(&trin,(char *)(name.c_str()),&val);
	gx = vtof(hdtype((char *)(name.c_str())),val);
            
	name="sx";
	gethdval(&trin,(char *)(name.c_str()),&val);
	x = gx - vtof(hdtype((char *)(name.c_str())),val);
            
	// adjust mute so that w=0 works
	if (mute_type) {
	  wm = MAX(dt/s,w);
	  float wtmp = mutefun((gx - s)/wm) * mutefun((tm - gx)/wm);
	  for (int i=0;i<nt;i++)
	    trout.data[i] = trin.data[i]*wtmp;
	}
	else {
	  wm=MAX(dt,w);
	  //      cerr<<"mute at offset="<<x<<" slope="<<s<<" zotime="<<tm<<" width="<<wm<<endl;
	  for (int i=0;i<nt;i++)
	    //	trout.data[i] = trin.data[i]*mutefun((t0+i*dt-s*fabs(x)-tm)/wm);
	    // this version: tm = time AFTER first sample
	    trout.data[i] = trin.data[i]*mutefun((i*dt+t0-s*fabs(x)-tm)/wm);
	}
      }
      catch (bad_cast) {
	RVLException e;
	e<<"Error: SEGYLinMute::operator()\n";
	e<<"input LDC not segytrace\n";
	throw e;
      }
    }
    
    
    void SEGYTaperMute::operator()(LocalDataContainer<float> & x,
				   LocalDataContainer<float> const & y) {
        
      try {
            
	segytrace & cpout = dynamic_cast<segytrace &>(x);
	segytrace const & cpin = dynamic_cast<segytrace const &>(y);
            
	segy & trout = cpout.getMetadata();
	segy const & trin = cpin.getMetadata();
            
	Value val;
	float gx;
	float sx;
	float x;
	float dt;
	float t0;
	float wm;
	float wt;
	int nt;
            
	string name="ns";
	gethdval(&trin,(char *)(name.c_str()),&val);
	nt=vtoi(hdtype((char *)(name.c_str())),val);
            
	name="dt";
	gethdval(&trin,(char *)(name.c_str()),&val);
	dt=0.001*vtof(hdtype((char *)(name.c_str())),val);
            
	name="delrt";
	gethdval(&trin,(char *)(name.c_str()),&val);
	t0=vtof(hdtype((char *)(name.c_str())),val);
            
	name="gx";
	gethdval(&trin,(char *)(name.c_str()),&val);
	gx = vtof(hdtype((char *)(name.c_str())),val);
            
	name="sx";
	gethdval(&trin,(char *)(name.c_str()),&val);
	sx = vtof(hdtype((char *)(name.c_str())),val);
	x = gx - sx; 
            
	int itw = (int) (tw/dt + 0.1);
	itw = MAX(1,itw);

	float sxt_wt = 1.0f;
	wt = MAX(1.0,sx_width);
	if (sx < sx_min) sxt_wt = cosfun2((sx_min-sx)/wt);
	if (sx > sx_max) sxt_wt = cosfun2((sx-sx_max)/wt);
	//cerr << " sxt_wt = " << sxt_wt << endl;
	// adjust mute so that w=0 works
	if (mute_type) {
	  wm = MAX(dt/s,w);
	  float wtmp = mutefun((gx - s)/wm) * mutefun((tm - gx)/wm);
	  wtmp = wtmp * sxt_wt;
	  wt=MAX(1.0,width);
	  float wttmp=1.0;
	  if (taper_type){ // offset
	    if (x < taper_min) wttmp = cosfun2((taper_min-x)/wt);
	    else if (x > taper_max) wttmp = cosfun2((x - taper_max)/wt);
#pragma ivdep                    
	    for (int i=0;i<nt-itw;i++)
	      trout.data[i] = trin.data[i]*wtmp*wttmp;
#pragma ivdep
	    for (int i=nt-itw;i<nt;i++)
	      trout.data[i] = trin.data[i]*wtmp*wttmp*cosfun(float(nt-i)/(itw+0.0f));

	  }
	  else { // taper geophone position
	    if (gx < taper_min) wttmp = cosfun2((taper_min-gx)/wt);
	    else if (gx > taper_max) wttmp = cosfun2((gx - taper_max)/wt);
#pragma ivdep                    
	    for (int i=0;i<nt-itw;i++)
	      trout.data[i] = trin.data[i]*wtmp*wttmp;
#pragma ivdep
	    for (int i=nt-itw;i<nt;i++)
	      trout.data[i] = trin.data[i]*wtmp*wttmp*cosfun(float(nt-i)/(itw+0.0f));
	  }
	}
	else {
	  wm=MAX(dt,w);
	  float wttmp=1.0;
	  wt=MAX(1.0,width);
	  if (taper_type){ // offset
	    if (x < taper_min) wttmp = cosfun2((taper_min-x)/wt);
	    else if (x > taper_max) wttmp = cosfun2((x - taper_max)/wt);
	    wttmp = wttmp * sxt_wt;
#pragma ivdep                    
	    for (int i=0;i<nt-itw;i++)
	      trout.data[i] = trin.data[i]*wttmp*mutefun((i*dt+t0-s*fabs(x)-tm)/wm);
#pragma ivdep
	    for (int i=nt-itw;i<nt;i++)
	      trout.data[i] = trin.data[i]*wttmp*mutefun((i*dt+t0-s*fabs(x)-tm)/wm)*cosfun(float(nt-i)/(itw+0.0f));
	  }
	  else { // taper geophone position
	    if (gx < taper_min) wttmp = cosfun2((taper_min-gx)/wt);
	    else if (gx > taper_max) wttmp = cosfun2((gx - taper_max)/wt);
	    wttmp = wttmp * sxt_wt;
#pragma ivdep
	    for (int i=0;i<nt-itw;i++)
	      trout.data[i] = trin.data[i]*wttmp*mutefun((i*dt+t0-s*fabs(x)-tm)/wm);
#pragma ivdep
	    for (int i=nt-itw;i<nt;i++)
	      trout.data[i] = trin.data[i]*wttmp*mutefun((i*dt+t0-s*fabs(x)-tm)/wm)*cosfun(float(nt-i)/(itw+0.0f));
	  }
	}
      }
      catch (bad_cast) {
	RVLException e;
	e<<"Error: SEGYTaperMute::operator()\n";
	e<<"input LDC not segytrace\n";
	throw e;
      }
    }
    
    void SEGYFwdInt::operator()(LocalDataContainer<float> & x,
                                LocalDataContainer<float> const & y) {
        
      try {
            
	segytrace & cpout = dynamic_cast<segytrace &>(x);
	segytrace const & cpin = dynamic_cast<segytrace const &>(y);
            
	segy & trout = cpout.getMetadata();
	segy const & trin = cpin.getMetadata();
            
	Value val;
	float dt;
	int nt;
            
	string name="ns";
	gethdval(&trin,(char *)(name.c_str()),&val);
	nt=vtoi(hdtype((char *)(name.c_str())),val);
	gethdval(&trout,(char *)(name.c_str()),&val);
	if (nt != vtoi(hdtype((char *)(name.c_str())),val)) {
	  RVLException e;
	  e<<"Error: SEGYFwdInt::operator()\n";
	  e<<"input, output incompatible\n";
	  throw e;
	}
	name="dt";
	gethdval(&trin,(char *)(name.c_str()),&val);
	dt=0.001*vtof(hdtype((char *)(name.c_str())),val);
            
	memcpy(trout.data,trin.data,nt*sizeof(float));
	trout.data[0]=0.0f;
	for (int j=0;j<nint;j++) {
#pragma ivdep
	  for (int i=1;i<nt;i++) trout.data[i]
				   = trout.data[i-1]+trin.data[i-1]*dt;
	}
      }
      catch (bad_cast) {
	RVLException e;
	e<<"Error: SEGYFwdInt::operator()\n";
	e<<"input LDC not segytrace\n";
	throw e;
      }
    }
    
    void SEGYAdjInt::operator()(LocalDataContainer<float> & x,
                                LocalDataContainer<float> const & y) {
        
      try {
            
	segytrace & cpout = dynamic_cast<segytrace &>(x);
	segytrace const & cpin = dynamic_cast<segytrace const &>(y);
            
	segy & trout = cpout.getMetadata();
	segy const & trin = cpin.getMetadata();
            
	Value val;
	float dt;
	int nt;
            
	string name="ns";
	gethdval(&trin,(char *)(name.c_str()),&val);
	nt=vtoi(hdtype((char *)(name.c_str())),val);
	gethdval(&trout,(char *)(name.c_str()),&val);
	if (nt != vtoi(hdtype((char *)(name.c_str())),val)) {
	  RVLException e;
	  e<<"Error: SEGYAdjInt::operator()\n";
	  e<<"input, output incompatible\n";
	  throw e;
	}
	name="dt";
	gethdval(&trin,(char *)(name.c_str()),&val);
	dt=0.001*vtof(hdtype((char *)(name.c_str())),val);
            
	memcpy(trout.data,trin.data,nt*sizeof(float));
	trout.data[nt-1]=0.0f;
	for (int j=0;j<nint;j++) {
#pragma ivdep
	  for (int i=nt-2;i>-1;i--) trout.data[i] 
				      = trout.data[i+1]+trin.data[i+1]*dt;
	}
      }
      catch (bad_cast) {
	RVLException e;
	e<<"Error: SEGYAdjInt::operator()\n";
	e<<"input LDC not segytrace\n";
	throw e;
      }
    }
    
    
  }