#include <iostream>
#include <cmath>
#include "TVector3.h"
#include "TExtrap.hh"
#include "Stntuple/obj/TStnTrack.hh"

//_____________________________________________________________________________
Int_t TExtrap::SetParameters(Double_t d0, Double_t phi, Double_t curv, 
		      Double_t z0, Double_t lambda)
{
  d = d0;
  p = phi;
  c = curv;
  z = z0;
  t = lambda;
  return FillStraight();
}

//_____________________________________________________________________________
Int_t TExtrap::SetParameters(TStnTrack* track)
{

  d = track->D0();
  p = track->Phi0();
  c = track->C0();
  z = track->Z0();
  t = track->Lam0();

  return FillStraight();
}

//_____________________________________________________________________________
Int_t TExtrap::SetParameters(TVector3* vertex, TVector3* pvec, Int_t charge)
{
  // this is not well-checked

  double pt = pvec->Pt();
  c = 0.0021159/pt;
  double R;
  if(c<1.0e-10) R = 1.0E10;  else R = 0.5/c;
  if(charge<0) c = -c;
  double x = vertex->X();
  double y = vertex->Y();
  double px = pvec->X();
  double py = pvec->Y();

  double xc,yc; // circle center
  if(charge>0){
    xc = x - R*py/pt;
    yc = y + R*px/pt;
  } else {
    xc = x + R*py/pt;
    yc = y - R*px/pt;
  }
  double h = sqrt(xc*xc+yc*yc);
  d = charge*(h-R);
  double px0,py0;
  if(charge>0) {
    px0 =  yc;
    py0 = -xc;
  } else {
    px0 =-yc;
    py0 = xc;
  }
  p = atan2(-py0,-px0) + M_PI;

  // z slope
  t = pvec->Z()/pvec->Pt();

  // z intercept
  double r = sqrt(x*x+y*y);
  double den = 1.0+2.0*c*d;
  double arg = r*r-d*d;
  double arg3 = arg/den;
  double k;
  if(fabs(c)<1.e-20) {
    k = t*sqrt(arg3);
  } else {
    k = (t/c)*asin(c*sqrt(arg3));
  }
  z = vertex->Z() - k;

  return FillStraight();

}

//_____________________________________________________________________________
Double_t TExtrap::PhiAtRInField(Double_t r)
{
  double den = 1.0 + 2.0*c*d;
  if(den<=0.0) return -999.0;
  double sinP = (r*c+(1.0+c*d)*d/r)/den;
  if(fabs(sinP)>=0.99999) return -999.0;
  double phi = p + asin(sinP);
  if(phi<0.0) phi += 2*M_PI;
  if(phi>2*M_PI) phi -= 2*M_PI;
  return phi;
}

//_____________________________________________________________________________
Double_t TExtrap::ZAtRInField(Double_t r)
{
  double den = 1.0+2.0*c*d;
  if(fabs(den)<1.e-20) return -999.0;
  double num = r*r - d*d;
  double arg = num/den;
  if(arg<1.e-20) return -999.0;
  double z1;
  if(fabs(c)<1.e-20) {
    z1 = t*sqrt(arg);
  } else {
    z1 = (t/c)*asin(c*sqrt(arg));
  }
  return z+z1;
}

//_____________________________________________________________________________
Double_t TExtrap::RAtZInField(Double_t zz)
{
  double den = 1.0+2.0*c*d;
  if(fabs(den)<1.e-20) return -999.0;
  double a1;
  if(fabs(c)>1.0e-20) {
    a1 = sin(c*(zz-z)/t)/c;
  } else {
    a1 = (zz-z)/t;
  }
  double a2 = a1*a1*den + d*d;
  if(a2<=0.0) return -999.0;
  return sqrt(a2);
}

//_____________________________________________________________________________
Double_t TExtrap::AlphaAtRInField(Double_t r)
{
  //angle swept through going from r=d to r=rcoil
  double den = 1.0+2.0*c*d;
  double sinDelPhi= 0.;
  double arg = r*r-d*d;
  if(fabs(arg) >1.0e-20) sinDelPhi = c*sqrt(arg/den);
  return 2.0*asin(sinDelPhi);
}


//_____________________________________________________________________________
Double_t TExtrap::PhiAtROutField(Double_t r)
{
  if(!outCoil) return -999.0;
  double arg = sd/r;
  if(fabs(arg)>0.9999) return -999.0;
  double phi = sp + asin(arg);
  if(phi<0.0) phi += 2*M_PI;
  if(phi>2*M_PI) phi -= 2*M_PI;
  return phi;
}

//_____________________________________________________________________________
Double_t TExtrap::ZAtROutField(Double_t r)
{
  if(!outCoil) return -999.0;
  double arg = r*r - sd*sd;
  if(arg<1.e-20) return -999.0;
  double dz = st*sqrt(arg);
  return sz+dz;
}

//_____________________________________________________________________________
Double_t TExtrap::RAtZOutField(Double_t zz)
{
  if(!outCoil) return -999.0;
  double a1 = (zz-sz)/st;
  double a2 = a1*a1 + d*d;
  return sqrt(a2);
}

//_____________________________________________________________________________
Int_t TExtrap::FillStraight()
{

  double r = rCoil;
  double phi = PhiAtRInField(r);
  if(phi<-998.0) {
    //cout << " didn't reach exit point " << endl;
    outCoil = false;
    return -1;
  }
  outCoil = true; // can go out to any r
  rexit = r;

  zexit = ZAtRInField(r); // where it exits field tube
  // check if it exits end of field
  if(zexit>zPes) {
    zexit = zPes;
    rexit = RAtZInField(zexit);
    r = rexit;
    phi = PhiAtRInField(r);
  }

  //phi angle swept through going from r=d to r=exiting field
  double alpha = AlphaAtRInField(r);

  // phi0 of the straight region now defined
  sp = p + alpha;
  if(sp<0.0)    sp += 2*M_PI;
  if(sp>2*M_PI) sp -= 2*M_PI;


  double adiff = phi-sp;
  if(adiff<-M_PI) adiff += 2.0*M_PI;
  if(adiff> M_PI) adiff -= 2.0*M_PI;

  // define d0 of straight region
  sd = r*sin(adiff);

  // lambda for straight region
  st = t;

  // z0 for straight region
  double srp = sqrt(r*r-sd*sd);
  sz = zexit - st*srp;

  //cout << " str sd,sp,sz,st " << sd << " " << sp << " " <<sz << " " <<st << endl;

  return 0;
}

//_____________________________________________________________________________
Int_t TExtrap::GetCesWedge()
{
  return GetWedge(24, rCes);
}

//_____________________________________________________________________________
Double_t TExtrap::GetCesLocalX(Int_t wedge)
{
  int w = (wedge<0?  GetWedge(24, rCes) : wedge);
  return GetLocalX(w, 24, rCes);
}

//_____________________________________________________________________________
Int_t TExtrap::GetCprWedge()
{
  return GetWedge(24, rCpr);
}

//_____________________________________________________________________________
Double_t TExtrap::GetCprLocalX(Int_t wedge)
{
  int w = (wedge<0?  GetWedge(24, rCpr) : wedge);
  return GetLocalX(w, 24, rCpr);
}

//_____________________________________________________________________________
Int_t TExtrap::GetWedge(Int_t nWedge, Double_t r)
{
  if(!outCoil) return -1;
  double phi = PhiAtROutField(r);
  int wedge = int( double(nWedge)*phi/(2.0*M_PI) );

  double xloc = GetLocalX(wedge, nWedge, r);
  double xloc2;
  if(xloc<0.0) {
    int wedge2 = (wedge==0? nWedge-1: wedge-1);
    xloc2 = GetLocalX(wedge2, nWedge, r);
    return (fabs(xloc)<fabs(xloc2)? wedge: wedge2);
  } else {
    int wedge2 = (wedge==nWedge-1? 0: wedge+1);
    xloc2 = GetLocalX(wedge2, nWedge, r);
    return (fabs(xloc)<fabs(xloc2)? wedge: wedge2);
  }
}

//_____________________________________________________________________________
Double_t TExtrap::GetLocalX(Int_t wedge, Int_t nWedge, Double_t r)
{
  if(!outCoil) return -999;
  double phiw = (double(wedge)+0.5)*2.0*M_PI/double(nWedge);
  double dphi = sp-phiw;
  return r*tan(dphi) + sd/cos(dphi);
}

//_____________________________________________________________________________
Double_t TExtrap::GetCesLocalZ(Int_t wedge)
{
  int w = (wedge<0?  GetWedge(24, rCes) : wedge);
  return GetLocalZ(w,24,rCes);
}

//_____________________________________________________________________________
Double_t TExtrap::GetCprLocalZ(Int_t wedge)
{
  int w = (wedge<0?  GetWedge(24, rCpr) : wedge);
  return GetLocalZ(w,24,rCpr);
}

//_____________________________________________________________________________
Int_t TExtrap::GetTofBar()
{
  double phi = PhiAtRInField(rTof);
  int bar = int(  216.0*phi/(2.0*M_PI) );
  return bar;
}

//_____________________________________________________________________________
Double_t TExtrap::GetTofLocalZ()
{
  return ZAtRInField(rTof);
}

//_____________________________________________________________________________
Double_t TExtrap::GetLocalZ(Int_t wedge, Int_t nWedge, Double_t r)
{
  // this ignores the wedge chamber flatness
  double xloc = GetLocalX(wedge, nWedge, r);
  double r2 = sqrt(xloc*xloc + r*r);
  return ZAtROutField(r2);
}

//_____________________________________________________________________________
Double_t TExtrap::PhiAtR(Double_t r)
{
  return ( r<rexit ? PhiAtRInField(r) :  PhiAtROutField(r) );
}

//_____________________________________________________________________________
Double_t TExtrap::ZAtR(Double_t r)
{
  return ( r<rexit ? ZAtRInField(r) :  ZAtROutField(r) );
}

//_____________________________________________________________________________
Double_t TExtrap::RAtZ(Double_t zz)
{
  bool useOut = (t>0.0 ? zz>zexit : zz<zexit );
  if(useOut) {
    return RAtZOutField(zz);
  } else {
    return RAtZInField(zz);
  }
}

//_____________________________________________________________________________
TVector3 TExtrap::PositionAtR(Double_t r)
{
  double phi = PhiAtR(r);
  double x = r*cos(phi);
  double y = r*sin(phi);
  double z = ZAtR(r);
  return TVector3(x,y,z);
}

//_____________________________________________________________________________
TVector3 TExtrap::MomentumAtR(Double_t r)
{

  double alpha;
  if( r>rexit ) {
    alpha = sp;
  } else {
    alpha = AlphaAtRInField(r)+p;
  }
  double pt =  0.0021159/(fabs(c)>1.0e-10? c : 1.0e-10);
  double px = pt*cos(alpha);
  double py = pt*sin(alpha);
  double pz = t*pt;
  return TVector3(px,py,pz);
}

//_____________________________________________________________________________
Double_t TExtrap::DetEtaAtShowerMax()
{
  if(IsCentral()) {
    double zces = ZAtROutField(rCes);
    double cott = zces/rCes;
    return log(sqrt(1+cott*cott)+cott);
  } else {
    double zz = (t>0? zPes: -zPes);
    double rpes = RAtZOutField(zz);
    double cott = zz/rpes;
    return log(sqrt(1+cott*cott)+cott);
  }
}


//_____________________________________________________________________________
Int_t TExtrap::GetPhiTower()
{
  if(IsCentral()) {
    return GetCesWedge();
  } else {
    return -1;
  }
}

//_____________________________________________________________________________
Int_t TExtrap::GetEtaTower()
{
  const float bound[11]={4.22, 24.16, 48.32, 72.48, 96.64,120.8,
			 144.97,169.12,193.28,217.45,245.96};
  if(IsCentral()) {
    double z = ZAtROutField(rCes);
    float az = fabs(z);
    int w = 0;
    while(w<9 && az>bound[w+1]) w++;
    if(z<0) w=-w-1;
    w += 26;    
    return w;
  } else {
    return -1;
  }
}