#include <ppc/i3ppc.h>
#include "ppc.h"

I3_MODULE(i3ppc);

using namespace std;

struct xtrPulses{
  double jit(){
    double t;
    do t=2*xppc::grnd(); while(t<-5 || t>25);
    return t*I3Units::ns;
  }

  double pts(double v){
    return -31.8*I3Units::ns*sqrt(1345*I3Units::V/v);
  }

  AfterPulseGeneratorPtr apg;
  LatePulseGeneratorPtr lpg;

  struct omi{
    double V, eff;
  };
  map<OMKey, omi> oms;
  map<OMKey, int> bad;
  double maxe;
  bool jpal_;

  void doJPAL(bool jpal){
    jpal_ = jpal;
  }

  void addOM(OMKey omkey, double V, double eff){
    if(!(V!=V || V+1==V || V<=0)){
      omi om;
      om.V=V;
      if(eff!=eff || eff+1==eff || eff<0) eff=1.0;
      om.eff=eff; if(maxe<eff) maxe=eff;
      if(bad.find(omkey)==bad.end()) oms[omkey]=om;
    }
  }

  void setr(const I3Context& c){
    shared_ptr<I3RandomService> r = c.Get<I3RandomServicePtr>("I3RandomService");
    if(!r) log_fatal("Failed to access RandomService");
    apg = AfterPulseGeneratorPtr(new AfterPulseGenerator(r));
    lpg = LatePulseGeneratorPtr(new LatePulseGenerator(r));
  }

  void addHit(OMKey & key, I3MCHitSeriesMap & hsm, I3MCHit & hit, double t){
    if(oms.find(key)==oms.end()) return;

    omi & om = oms[key];
    if(om.eff<maxe*xppc::xrnd()) return;

    double v=om.V, rnd=xppc::xrnd();
    I3MCHitSeries & hs = hsm[key];

    const double pre=0.007;
    const double late=0.035;
    const double after=0.0593;

    if(jpal_) { // Simulate Pre After Late pulses
      if(rnd<pre){
        hit.SetHitSource(I3MCHit::PRE_PULSE);
        t+=pts(v)+jit();
      }
      else if(rnd<pre+late){
        hit.SetHitSource(I3MCHit::ELASTIC_LATE_PULSE);
        t=lpg->GenerateSingleLatePulse(t, v);
      }
      else{
        hit.SetHitSource(I3MCHit::SPE);
        t+=jit();
      }
      hit.SetTime(t);
      hs.push_back(hit);

      if(xppc::xrnd()<after){
         vector<I3MCHit> aps;
         apg->GenerateAfterPulses(aps, hit);
         for(vector<I3MCHit>::const_iterator i=aps.begin(); i!=aps.end(); ++i) hs.push_back(*i);
      }
    } 
    else { // PAL happens later in simulation
        hit.SetHitSource(I3MCHit::SPE);
        hit.SetTime(t);
        hs.push_back(hit);
    }
  }
} pz;

i3ppc::i3ppc(const I3Context& ctx) : I3Module(ctx),
  ini(false), 
  verbose(false), 
  pal_(false), 
  gpu(0), 
  nph(0)
{
  AddParameter ("gpu", "GPU to use", gpu);
  AddParameter ("bad", "DOMs not to use", bad);
  AddParameter ("fla", "Flasher position", fla);
  AddParameter ("nph", "Number of photons", nph);
  AddParameter ("JPALpulses", "Simulate Jitter and Pre,After,Late pulses", pal_);

  AddOutBox("OutBox");

  xppc::start();

  Radius=0, Top=0, Bottom=0, Padding=300;
}

i3ppc::~i3ppc(){
  xppc::stop();
}

void i3ppc::Configure(){
  GetParameter ("gpu", gpu);
  GetParameter ("bad", bad);
  GetParameter ("fla", fla);
  GetParameter ("nph", nph);
  GetParameter ("JPALpulses", pal_);
  xppc::initialize();
  xppc::choose(gpu);
  fb=0, fe=0;
  pz.setr(context_);
  pz.doJPAL(pal_);
  for(vector<OMKey>::const_iterator i=bad.begin(); i!=bad.end(); ++i) pz.bad[*i]++;
}

void i3ppc::DAQ(I3FramePtr frame){
  log_trace("Entering i3ppc::DAQ()");
  frames.push_back(frame);

  if(!ini){
    if(verbose) cout<<"Geometry frame"<<endl;
    I3GeometryConstPtr geo = frame->Get<I3GeometryConstPtr>();
    I3CalibrationConstPtr calib = frame->Get<I3CalibrationConstPtr>();
    I3DetectorStatusConstPtr status = frame->Get<I3DetectorStatusConstPtr>();
    if(verbose){
      if(!geo) cout<<"    --> empty"<<endl;
      else cout<<"    --> contains "<<geo->omgeo.size()<<" entries"<<endl;
    }
    if(geo) if(geo->omgeo.size()>0){
      pz.maxe=0;

      for(I3OMGeoMap::const_iterator it=geo->omgeo.begin(); it!=geo->omgeo.end(); ++it){
	OMKey omkey=it->first;
	xppc::OM om;
	om.str=omkey.GetString();
	om.dom=omkey.GetOM();
	const I3Position& pos = (it->second).position;
	om.r[0]=pos.GetX()/I3Units::m;
	om.r[1]=pos.GetY()/I3Units::m;
	om.r[2]=pos.GetZ()/I3Units::m;
	xppc::i3oms.push_back(om);

	if(status){
	  map<OMKey, I3DOMStatus>::const_iterator om_stat = status->domStatus.find(omkey);
	  if(om_stat!=status->domStatus.end()){
	    double eff=1.0;
	    if(calib){
	      map<OMKey, I3DOMCalibration>::const_iterator om_cal = calib->domCal.find(omkey);
	      if(om_cal!=calib->domCal.end()) eff=om_cal->second.GetRelativeDomEff();
	    }
	    pz.addOM(omkey, om_stat->second.pmtHV, eff);
	  }
	}

	double R=sqrt(om.r[0]*om.r[0]+om.r[1]*om.r[1]); if(R>Radius) Radius=R;
	if(om.r[2]>Top) Top=om.r[2]; else if(om.r[2]<Bottom) Bottom=om.r[2];
      }

      xppc::igeo(pz.maxe);
      if(nph>0){
	xppc::flini(fla.GetString(), fla.GetOM());
	xppc::sett(0, 0, 1, make_pair(-1, -1), fe++);
      }
      else xppc::ini(0);
      ini=true;
    }
  }

  if(nph>0){
    xppc::flone((unsigned long long)llroundf(nph)); fb=0;
  }
  else{
    for(I3Frame::const_iterator iter=frame->begin(); iter!=frame->end(); iter++){
      I3MMCTrackListConstPtr mmclist = dynamic_pointer_cast<const I3MMCTrackList>(iter->second);
      if(mmclist){
	if(verbose) cout<<"  --> This is an I3MMCTrackList: "<<iter->first<<endl;
	for(I3MMCTrackList::const_iterator it=mmclist->begin(); it!=mmclist->end(); ++it){
	  const I3MMCTrack& m=*it;
	  const I3Particle& p=it->GetI3Particle();
	  int minor = p.GetMinorID();
	  unsigned long long major = p.GetMajorID();
	  if(verbose) cout<<"      MMCTrack info added for particle with id: ("<<major<<","<<minor<<")"<<endl;
	  i3mmctracks[make_pair(minor, major)]=&m;
	}
      }
    }

    xppc::eini();
    bool trees=false, lists=false;
    for(I3Frame::const_iterator iter=frame->begin(); iter!=frame->end(); iter++){
      I3MCTreeConstPtr mctree = dynamic_pointer_cast<const I3MCTree>(iter->second);
      if(mctree){
	if(verbose) cout<<"  --> This is an I3MCTree: "<<iter->first<<endl;
	if(lists) continue; else trees=true;
      }
      else{
	I3MCListConstPtr mclist = dynamic_pointer_cast<const I3MCList>(iter->second);
	if(mclist){
	  if(verbose) cout<<"  --> This is an I3MCList: "<<iter->first<<endl;
	  if(trees) continue; else lists=true;
	  mctree = I3MCTreeUtils::ListToTree(mclist);
	  if(mctree) if(verbose) cout<<"  --> Converted to I3MCTree: "<<iter->first<<endl;
	}
      }
      if(mctree){
	int num=0;
	for(I3MCTree::sibling_iterator it=mctree->begin(); it!=mctree->end(); ++it, num++){
	  if(verbose) cout<<"extracting a particle "<<it->GetType()<<endl;
	  pparticle(mctree, it, num);
	}
      }
    }

    i3mmctracks=map<pair<int, unsigned long long>, const I3MMCTrack *>();
    fe++;
  }

  pushframe();
}

void i3ppc::popframes(int fx){
  for(; fb<fx; fb++){
    bool empty = true;
    if(mchits.find(fb)==mchits.end()) mchits[fb]=I3MCHitSeriesMapPtr(new I3MCHitSeriesMap);
    else empty = mchits[fb]->empty();

    if(mchits[fb]) frames.front()->Put("MCHitSeriesMap", mchits[fb]);
    mchits.erase(mchits.find(fb));

    if(!empty){
      log_debug("Pushing the frame");
      PushFrame(frames.front(),"OutBox");
    }
    frames.pop_front();
  }
}

void i3ppc::pushframe(){
  int fc=0;
  for(unsigned int i=0; i<xppc::hitz.size(); i++){
    xppc::ihit & h = xppc::hitz[i];
    OMKey key(h.omkey.str, h.omkey.dom);

    fc=h.track.frame;
    if(mchits.find(fc)==mchits.end()) mchits[fc]=I3MCHitSeriesMapPtr(new I3MCHitSeriesMap);

    I3MCHit hit(h.track.second, h.track.first);
    hit.SetHitID(i); hit.SetWeight(1.0);
    pz.addHit(key, *mchits[fc], hit, h.time*I3Units::ns);
  }
  xppc::efin();
  popframes(nph>0?fe:fc);
}

void i3ppc::Finish(){
  if(ini){
    if(nph<=0){
      xppc::eout();
      pushframe();
      popframes(fe);
    }
    Flush();
    xppc::fin();
  }
}

double i3ppc::res(double x, double mod){
  return x-mod*int(x/mod);
}

void i3ppc::pparticle(I3MCTreeConstPtr tree, I3MCTree::iterator sub, int n){
  static int level=0;

  const I3Particle& p = *sub;
  int type=p.GetType(), minor=p.GetMinorID();
  unsigned long long major=p.GetMajorID();
  pair<int, unsigned long long> id=make_pair(minor, major);

  if(verbose){
    for(int i=0; i<level; i++) cout<<"    ";
    cout<<"      --> particle # "<<n<<" of type "<<type<<" id: ("<<major<<","<<minor<<")"<<endl;
  }

  int ptype = type==I3Particle::MuPlus || type==I3Particle::MuMinus ? -1 :
    type==I3Particle::DeltaE || type==I3Particle::Brems || type==I3Particle::PairProd ||
    type==I3Particle::EPlus || type==I3Particle::EMinus ? 0 :
    type==I3Particle::NuclInt || type==I3Particle::Hadrons ? 1 : -2;

  if(ptype==-1){
    const double sol=0.299792458;
    double E0=p.GetEnergy()/I3Units::GeV;
    double t0=p.GetTime()/I3Units::ns;
    double ll=p.GetLength()/I3Units::m;

    double Ei, Ef, ti, tf;

    map<pair<int, unsigned long long>, const I3MMCTrack *>::const_iterator mmctrack=i3mmctracks.find(id);
    if(mmctrack!=i3mmctracks.end()){
      const I3MMCTrack * i3mmctrack = mmctrack->second;
      ti=i3mmctrack->GetTi()/I3Units::ns;
      Ei=i3mmctrack->GetEi()/I3Units::GeV;
      tf=i3mmctrack->GetTf()/I3Units::ns;
      Ef=i3mmctrack->GetEf()/I3Units::GeV;
      if(Ei==0){
	ti=t0;
	Ei=E0;
      }
      if(Ef<0){
	tf=t0+ll/sol;
	Ef=0;
      }
      if(Ei>0){
	if(verbose) cout<<"extracting a muon Ei="<<Ei<<" "<<tree->number_of_children(sub)<<" secondaries"<<endl;

	double th=180-p.GetZenith()/I3Units::degree;
	double ph=res(180+p.GetAzimuth()/I3Units::degree, 360);

	double sinth=sin(th*M_PI/180);
	double costh=cos(th*M_PI/180);
	double sinph=sin(ph*M_PI/180);
	double cosph=cos(ph*M_PI/180);

	double nx=sinth*cosph;
	double ny=sinth*sinph;
	double nz=costh;

	double totE=0;
	for(I3MCTree::sibling_iterator it=tree->begin(sub); it!=tree->end(sub); ++it){
	  double t=it->GetTime()/I3Units::ns;
	  if(t>=ti && t<=tf) totE+=it->GetEnergy()/I3Units::GeV;
	}
	double sl0=tf>ti?((Ef-Ei+totE)/(ti-tf)):0;
	double slmax=(0.21+8.8e-3*log(Ei)/log(10.))*sol;  // for ice only at Ecut=500 MeV
	double sl=min(sl0, slmax);

	double x=p.GetX()/I3Units::m;
	double y=p.GetY()/I3Units::m;
	double z=p.GetZ()/I3Units::m;
	double t=ti;
	double l=sol*(tf-ti);
	double E=Ei;
	double new_t=tf;

	for(I3MCTree::sibling_iterator it=tree->begin(sub); it!=tree->end(sub); ++it){
	  new_t=it->GetTime()/I3Units::ns;
	  if(new_t>=ti && new_t<=tf){
	    l=sol*(new_t-t);

	    if(verbose) cout<<"muon at z="<<z<<" length="<<l<<" E="<<E<<" t="<<t<<endl;
	    if(ptype==-1){
	      xppc::sett(nx, ny, nz, id, fe);
	      xppc::addp(x, y, z, t, E, (float) l);
	    }

	    x=it->GetX()/I3Units::m;
	    y=it->GetY()/I3Units::m;
	    z=it->GetZ()/I3Units::m;
	    double e=it->GetEnergy()/I3Units::GeV;

	    E-=e+sl*max(new_t-t, 0.);
	    t=new_t;

	    {
	      int type=it->GetType();
	      int ptype = type==I3Particle::DeltaE ||
		type==I3Particle::Brems || type==I3Particle::PairProd ||
		type==I3Particle::EPlus || type==I3Particle::EMinus ? 0 :
		type==I3Particle::NuclInt || type==I3Particle::Hadrons ? 1 : -3;

	      if(ptype>=0){
		xppc::sett(nx, ny, nz, make_pair(it->GetMinorID(), it->GetMajorID()), fe);
		if(verbose) cout<<(ptype==0?"e-m":"hadr")<<" cascade at z="<<z<<" E="<<e<<" t="<<t<<endl;
		xppc::addp(x, y, z, t, e, ptype);
	      }
	    }
	  }
	}
	l=sol*(tf-t);
	if(verbose) cout<<"muon at z="<<z<<" length="<<l<<" E="<<E<<" t="<<t<<endl;
	if(ptype==-1){
	  xppc::sett(nx, ny, nz, id, fe);
	  xppc::addp(x, y, z, t, E, (float) l);
	}
      }
    }
    else{
      ti=t0;
      Ei=E0;
      tf=t0+ll/sol;
      Ef=0;
    }
  }
  else if(ptype>=0){
    double x=p.GetX()/I3Units::m;
    double y=p.GetY()/I3Units::m;
    double z=p.GetZ()/I3Units::m;

    if(sqrt(x*x+y*y)<Radius+Padding && Bottom-Padding<z && z<Top+Padding){
      double th=180-p.GetZenith()/I3Units::degree;
      double ph=res(180+p.GetAzimuth()/I3Units::degree, 360);

      double sinth=sin(th*M_PI/180);
      double costh=cos(th*M_PI/180);
      double sinph=sin(ph*M_PI/180);
      double cosph=cos(ph*M_PI/180);

      double nx=sinth*cosph;
      double ny=sinth*sinph;
      double nz=costh;

      xppc::sett(nx, ny, nz, id, fe);

      double t=p.GetTime()/I3Units::ns;
      double e=p.GetEnergy()/I3Units::GeV;

      if(verbose) cout<<(ptype==0?"e-m":"hadr")<<" cascade at z="<<z<<" E="<<e<<" t="<<t<<endl;
      xppc::addp(x, y, z, t, e, ptype);
    }
  }
  else{
    level++;
    int num=0;
    for(I3MCTree::sibling_iterator it=tree->begin(sub); it!=tree->end(sub); ++it, num++) pparticle(tree, it, num);
    level--;
  }

}