///////////////////////////////////////////////////////////////////////////////
//
//  MuoSegmentReco.cpp
//
//   4/30/99 Yasuda Switched to new framework macros
//   11/07   Hedin change mixed region from .3x.3 to .2x.3
//   5/08 Hedin change back to original as doesn't seem to matter
///////////////////////////////////////////////////////////////////////////////

#include "framework/Registry.hpp"

#include "muo_segmentreco/MuoSegmentReco.hpp"
#include "muon_geometry/base/MuoGeometer.hpp"

#include "muo_segmentreco/MuoSegmentChunk.hpp"

#include "muon_index/MuoSectionIndex.hpp"
#include "geometry_system/components/GenTrapezoid.hpp"

#include "edm/LinkPtr.hpp"

#include "muo_hitreco/MDTHitChunk.hpp"
#include "muo_hitreco/MSCHitChunk.hpp"
#include "muo_hitreco/PDTHitChunk.hpp"

#include "muo_hit/MDTGeometryHit.hpp"
#include "muo_hit/MSCGeometryHit.hpp"
#include "muo_hit/PDTGeometryHit.hpp"

#include "muo_util/Log.hpp"

#include "muo_segment/SegmentChunk.hpp"
#include "muo_segment/SegmentBuilder.hpp"
#include "muo_segmentlinkedlist/LinkedListSegmentBuilder.hpp"
#include "muo_segmentcombi/CombinatorialSegmentBuilder.hpp"

#include "kinem_util/AnglesUtil.hpp"
#include "muo_util/MCFlag.hpp"
#include "identifiers/CollisionID.hpp"
#include "muo_util/RunNumber.hpp"
#include "prod_history/HistorySelector.hpp"
#include "muo_util/MuonCalibrationConstants.hpp"
#include "muon_PDT_calibration/MuonPdtRecoCalibrater.hpp"
#include "muon_geometry/base/MuoXForm.hpp"

//using namespace Muon;
using namespace dgs;
using namespace fwk;
using fwk::HistorySelector;
 
namespace Muon {

FWK_REGISTRY_IMPL(MuoSegmentReco, "$Name:  $")

// Constructor
MuoSegmentReco::MuoSegmentReco(fwk::Context* context): 
  fwk::Package(context), 
  fwk::Process(context)
{
  _segmentBuilder = 0;
  _scintBuilder = 0;

  edm::RCP rcp = packageRCP();
  std::string algName = rcp.getString( "SegmentAlgorithm" );
  if ( algName == "Combinatorial" )
  {
    _segmentBuilder = new CombinatorialSegmentBuilder(rcp.getRCP("Combi"));
  } 
  else if( algName == "LinkedList" ) 
    {
      _segmentBuilder = new LinkedListSegmentBuilder(rcp.getRCP("LL"));
    } 
  else
    {
      _segmentBuilder = 0;
      std::cout << "Invalid muon segment algorithm" << std::endl;
    } 
  _scintBuilder = new MuoScintSegmentBuilder(rcp);

  _rcpID = rcp.getRCPID();
  _debugLevel = rcp.getInt("DebugLevel");
  _deleteOldChunks = rcp.getBool("deleteOldChunks");

  //=== if one wants to apply the RCP calib
  bool use_rcp_calib = rcp.getBool("use_rcp_calib");
  _use_PDTdb_calib = rcp.getBool("use_PDTdb_calib");
  if(!_use_PDTdb_calib)
    {
      MuonCalibrationConstants::Initialize(rcp.getRCP("Constants"));
    }
  else
    {
      MuonCalibrationConstants::InitializeDB();
    }
  // PDT pads fit
  _use_pads = rcp.getBool("use_Pads");
  _nsol = rcp.getInt("nSolutions");
  _chi2oldpos = rcp.getFloat("chi2Oldpos");
  _minDecks = rcp.getInt("minDecks");

}

MuoSegmentReco::~MuoSegmentReco()
  //{}
  {
  delete _segmentBuilder;
  delete _scintBuilder;
  } 
  

////////////////////////////////////////////processEvent method
//
fwk::Result MuoSegmentReco::processEvent(edm::Event& event)
{  
   if (_debugLevel >= 1) {
    std::cout << "Starting processing segments" << std::endl;
  }

   //
   // Check if the event is Monte Carlo or data
   //
   bool flag = HistorySelector::is_monte_carlo(event);
   MCFlag::set(flag);
  
   
   //=== Access run number and store it in Muon::RunNumber ===
   int run_number = event.collisionID().runNumber();
   RunNumber::set(run_number);


   if(_use_PDTdb_calib && !MCFlag::test())
     {
       MuonPdtCalibrater* calibrater = MuonPdtRecoCalibrater::get_instance();
       if(calibrater->updated())
	 {
	   MuonCalibrationConstants::UpdateDB();
	   calibrater->reset_updated();
	 }
     }

 
  MuoGeometer * geometer=MuoGeometer::get_instance();
  if( 0 == geometer ) return fwk::Result::success;
  
  const edm::TKey<MDTHitChunk> mdthitkey;
  edm::THandle<MDTHitChunk> mdtchunk = mdthitkey.find(event);
  
  const edm::TKey<MSCHitChunk> mschitkey;
  edm::THandle<MSCHitChunk> mscchunk = mschitkey.find(event);
  
  const edm::TKey<PDTHitChunk> pdthitkey;
  edm::THandle<PDTHitChunk> pdtchunk = pdthitkey.find(event);
  
  if( !mdtchunk.isValid() ) {
    Log::log()(ELwarning2, "No MDTChunk found for segment reconstruction") << endmsg;
    return Result::success;
  }
  
  if( !mscchunk.isValid() ) {
    Log::log()(ELwarning2, "No MSCChunk found for segment reconstruction") << endmsg;
    return Result::success;
  }
  
  if( !pdtchunk.isValid() ) {
    Log::log()(ELwarning2, "No PDTChunk found for segment reconstruction") << endmsg;
    return Result::success;
  }
  
  // build the list of parent chunks
  std::list<edm::ChunkID> parents;
  parents.push_back(mdtchunk->chunkID());
  parents.push_back(pdtchunk->chunkID());
  parents.push_back(mscchunk->chunkID());
  // Now create an empty output chunk
  MuoSegmentChunk *muosegchunk = new MuoSegmentChunk (parents, _rcpID);
  std::auto_ptr<MuoSegmentChunk> segchunk (muosegchunk);
  
  if( _deleteOldChunks ) {
    // Delete all previous MuoSegmentChunks
    edm::TKey<MuoSegmentChunk> muosegkey;
    std::list<edm::THandle<MuoSegmentChunk> > muosegchunks = muosegkey.findAll(event);    
    
    for( std::list<edm::THandle<MuoSegmentChunk> >::iterator mciter = muosegchunks.begin();
	 mciter != muosegchunks.end();
	 ++mciter ) {
      event.deleteChunk( (*mciter)->chunkID() );
    }    

    // Delete all previous SegmentChunks
    edm::TKey<SegmentChunk> segkey;
    std::list<edm::THandle<SegmentChunk> > segchunks = segkey.findAll(event);

    for( std::list<edm::THandle<SegmentChunk> >::iterator citer = segchunks.begin();
	 citer != segchunks.end();
	 ++citer ) {
      event.deleteChunk( (*citer)->chunkID() );
    }
  }

  
       
  
  // Get the processed hits and build segements 
  // Get the hits out of the chunk
  const MDTHitCollection &mdtcollection = mdtchunk->getHits();
  const MSCHitCollection &msccollection = mscchunk->getHits();
  const PDTHitCollection &pdtcollection = pdtchunk->getHits();
    
  if( _segmentBuilder != 0 ) {  
    SegmentCollection segments = _segmentBuilder->getSegments(pdtcollection,
							      mdtcollection,
							      msccollection);   

    // Here follows the building of scintillator-only segments
    _scintBuilder->findSegments(segments, msccollection);

    // fix the region of mixed segments
    fixRegion(segments);

    // use PDT pads to get coordinate along wire
    if (_use_pads) fit_pads(segments);

    // Bridge to old interface
    MuoSegmentChunk *muosegmentchunk = new MuoSegmentChunk(parents, _rcpID);
    muosegmentchunk->setSegmentCollection(segments);

    std::auto_ptr<MuoSegmentChunk> muosegmentchunkptr(muosegmentchunk);
    edm::insertChunk(event, muosegmentchunkptr);

    // Finish up the links inside the segments
    // Don't do this for the combi algorithm!
    //if( packageRCP().getString( "SegmentAlgorithm" ) == "LinkedList" ) {
      for( SegmentCollection::iterator iter = segments.begin();
	   iter != segments.end();
	   ++iter ) {
	iter->updatePDTLinks(pdtchunk->chunkID(), pdtcollection);
	iter->updateMDTLinks(mdtchunk->chunkID(), mdtcollection);
	iter->updateMSCLinks(mscchunk->chunkID(), msccollection);
      }
      //}

    // Now we have to put in the new SegmentChunk in the event
    SegmentChunk *segmentchunk = new SegmentChunk(parents, _rcpID);
    segmentchunk->setSegmentCollection(segments);

    std::auto_ptr<SegmentChunk> segmentchunkptr(segmentchunk);
    edm::insertChunk(event, segmentchunkptr);
  }

 
  return Result::success;
}

void MuoSegmentReco::fixRegion(SegmentCollection &segments) {

// After all segments have been found, we adapt the region of 
// some segments to match that of mixed region segments

// This code: change PDT region if match to mix
    int cntr = -1;
    for (SegmentCollection::iterator seg1 = segments.begin();
         seg1 != segments.end(); ++seg1) {
      cntr++;
      int npdt1 = seg1->getPDTLinks().size();
      int nmdt1 = seg1->getMDTLinks().size();
      int nmsc1 = seg1->getMSCLinks().size();
      int layer1 = seg1->getLayer();
      int region1 = seg1->getRegion();
      bool isA1 = (layer1 == 0);
      if (region1==0 && nmdt1==0 && ((npdt1>0)||(npdt1==0&&nmsc1>0))) {
        double eta1 = seg1->eta();
        SpacePoint pos1 = seg1->getPosition();
        double etap1 = kinem::eta(pos1.theta());
        double phi1 = seg1->phi();
        int ctr2 = -1;
        int nmixfwd = 0;
        int npuurbar = 0;
        int region_of_mix = -1;
        for (SegmentCollection::iterator seg2 = segments.begin();
             seg2 != segments.end(); ++seg2) {
          ctr2++;
          int npdt2 = seg2->getPDTLinks().size();
          int nmdt2 = seg2->getMDTLinks().size();
          int nmsc2 = seg2->getMSCLinks().size();
          double eta2 = seg2->eta();
          SpacePoint pos2 = seg2->getPosition();
          double etap2 = kinem::eta(pos2.theta());
          double phi2 = seg2->phi();
          int region2 = seg2->getRegion();
          int layer2 = seg2->getLayer();
          bool isA2 = (layer2 == 0);
          //double deta = abs(eta1-eta2);
          double deta = abs(etap1-etap2);
          double dphi = abs(phi1-phi2);
          if (dphi > kinem::PI) {dphi = kinem::TWOPI - dphi;}
          if (((isA1&&!isA2)||(isA2&&!isA1))&&(dphi<0.3)&&(deta<0.3)) {
//          if (((isA1&&!isA2)||(isA2&&!isA1))&&(dphi<0.3)&&(deta<0.2)) {
            if (npdt2 > 0 && nmdt2 > 0) {
              nmixfwd++;
              region_of_mix = region2;
            }
            if (npdt2 > 0 && nmdt2 == 0) {
              npuurbar++;
            }
          }
        }
        if (nmixfwd>0) {
          int octant1 = seg1->getOctant();
          double dr_angle = seg1->dr_angle();
          double x_error = seg1->x_error();
          double y_error = seg1->y_error();
          double z_error = seg1->z_error();
          double xenew = 0;
          double yenew = 0;
          double zenew = 0;
          if ((octant1==0)||(octant1==3)||(octant1==4)||(octant1==7)) {
            yenew = y_error;
            xenew = abs(z_error*tan(dr_angle));
          }
          if ((octant1==1)||(octant1==2)||(octant1==5)||(octant1==6)) {
            xenew = x_error;
            yenew = abs(z_error*tan(dr_angle));
          }
          seg1->setRegion(region_of_mix);
          seg1->set_x_error(xenew);
          seg1->set_y_error(yenew);
          seg1->set_z_error(zenew);
          int new_region = seg1->getRegion();
          x_error = seg1->x_error();
          y_error = seg1->y_error();
          z_error = seg1->z_error();
        }
      }
    }

}

void MuoSegmentReco::fit_pads(SegmentCollection &segments) {

 for( SegmentCollection::iterator segiter = segments.begin();
                                  segiter != segments.end(); ++segiter ) {

   double newsegpos = -999;
   double rms = 999;
   int ndeck=0;
   bool acceptpads = false;
   int nsol = _nsol;
   float chi2max = _chi2oldpos;

   int oct = segiter->getOctant();

   double segpos = -999;
   double segerror = 999.;

   //definte pad positions in terms of a zig-zag function:
   double wl = 60.96; //wavelength in cm
   double amp = 0.625; //amplitude (in terms of the pad difference);
   double shift = 0.09; //a global offset in the pad ratios.
   double slope = 4*amp / wl; // slope of zig-zag.
   double off[4];
   double posmax=0, posmin=0;//use to ensure the solutions lie inside the PDT!

   // these octants have somewhat different parameters
   if (oct==2 || oct==5 || oct==6) {
     amp = 0.7;
     shift = 0.2;
   }

   // segpos is the coordinate we want to measure (along the wire)
   // planepos is the position of the plane at which segment pars are given
   double planepos = 999.;
   if(oct==0 || oct==7){
     //define the offsets of the pad pattern in each deck:
     off[0] = 5.5;
     off[1] = off[3] = 25.82;
     off[2] = 15.66;
     segpos = segiter->y();
     segerror = segiter->y_error();
     posmax = 310;
     posmin = -310;
     planepos = segiter->x();
   }
   if(oct==3 || oct==4){
     //define the offsets of the pad pattern in each deck:
     off[0] = 3.82;
     off[1] = off[3] = 24.5;
     off[2] = 14.65;
     segpos = segiter->y();
     segerror = segiter->y_error();
     posmax = 310;
     posmin = -310;
     planepos = segiter->x();
   }
   if(oct==1){
     //define the offsets of the pad pattern in each deck:
     off[0] = 23.0;
     off[1] = off[3] = 3.5;
     off[2] = 13.5;
     segpos = segiter->x();
     segerror = segiter->x_error();
     posmax = 310;
     posmin = -10;
     planepos = segiter->y();
   }
   if(oct==2){
     //define the offsets of the pad pattern in each deck:
     off[0] = 36.3;
     off[1] = off[3] = 56.2;
     off[2] = 46.1;
     segpos = segiter->x();
     segerror = segiter->x_error();
     posmax = 10;
     posmin = -310;
     planepos = segiter->y();
   }
   if(oct==5){
     //define the offsets of the pad pattern in each deck:
     off[0] = 42.6;
     off[1] = off[3] = 32.4;
     off[2] = 52.2;
     segpos = segiter->x();
     segerror = segiter->x_error();
     posmax = 10;
     posmin = -310;
     planepos = segiter->y();
   }
   if(oct==6){
     //define the offsets of the pad pattern in each deck:
     off[0] = 19.2;
     off[1] = off[3] = 29.4;
     off[2] = 9.1;
     segpos = segiter->x();
     segerror = segiter->x_error();
     posmax = 310;
     posmin = -10;
     planepos = segiter->y();
   }
   if (abs(planepos)<1.) continue;

   //an array to store the best solutions in the 4 decks:
   double pbest[4][4];
   //this array stores the solutions modified for track angle
   double qbest[4][4];
   //this array stores the wire positions for each plane hit
   double wbest[4];
   for(int i=0; i<4; i++) {
     wbest[i]=-999;
     for(int j=0; j<4; j++) {
       pbest[i][j]=-999;
       qbest[i][j]=-999;
     }
   }

   // loop over the PDT hits
   const PDTLinkCollection& pdtlinks = segiter->getPDTLinks();
   PDTLinkCollection::const_iterator pdtlink_iter;
   for(pdtlink_iter=pdtlinks.begin(); pdtlink_iter!=pdtlinks.end(); 
             pdtlink_iter++)
     {
       edm::LinkPtr<PDTHitChunk, PDTHit> pdt_ptr(*pdtlink_iter);
       if( pdt_ptr.isValid() ) {
	
	const MuoIndex& index = pdt_ptr->getMuoIndex();

	//require a PDT hit (just in case!):
	if(index.type()!=0 || index.region()!=0 || index.layer()!=0) continue;
	
	//get the cell and deck numbers:
	int eta = index.eta();
	int plane = index.plane();

        // get wire coordinates
        double a=0;
        PDTHit p = *pdt_ptr;
        PDTGeometryHit* geohit = new PDTGeometryHit(&p,a);
        SpacePoint xyz = geohit->getWirePosition();
        delete geohit;
        if (oct==0 || oct==3 || oct==4 || oct==7) {
          wbest[plane] = xyz.x();
        } else {
          wbest[plane] = xyz.y();
        }
        if (abs(wbest[plane])<1.) wbest[plane]=1.;
	
	//determine if we should look for pads 1&2 or 3&4:
	int even;
	if( oct==0 || oct==1 || oct==5 || oct==7 ) {
	  if(eta%2==1) even=0;
	  else even = 1;
	}
	else {
	  if(eta%2==1) even=1;
	  else even = 0;
	}
	
	double norm=0; // the normalised pad difference
	int p1=0,p2=0; //the adc's on each pad
	if(even && pdt_ptr->getPadStatus(0)==0 && pdt_ptr->getPadStatus(1)==0){
	  p1=pdt_ptr->getPad(0);
	  p2=pdt_ptr->getPad(1);
	}
	if(even==0 && pdt_ptr->getPadStatus(2)==0 && 
                                                  pdt_ptr->getPadStatus(3)==0){
	  p1=pdt_ptr->getPad(2);
	  p2=pdt_ptr->getPad(3);
	}
	
	if(p1>0 && p2>0 && p1<940 && p2<940) {
	  norm = (double) (p1-p2) / (p1+p2);
	}
	else continue; //i.e. if pads not readout out correctly, or saturated.

        // protect against too large, or too small, norm
        if (norm > (amp-shift)) norm = amp - shift;
        if (norm < (-amp-shift)) norm = -amp - shift;

	double res =0;	
	//now we calculate all the posible solutions along the zig-zag which 
        //match the pad ratio.
	//start by calculating the solutions in one up and one down part of the 
        //zig-zag
	//assuming 0,0 is the bottom left corner of the first (up) part of the 
        //zig-zag
	double xpos = norm;
	xpos += shift + amp; //the total shift in the 'y' (pad ratio) direction.
	xpos  /= slope; // convert this to the position using the y/x slope
	double xpos2 = wl;
	xpos2 -= xpos; //the mirror-image solution on the downward part of the 
                       //zig-zag
	
	xpos += off[plane]; //add the pad offest in the 'x' direction
	xpos2 += off[plane]; //add the pad offest in the 'x' direction
	
	//now calculate all 11 possible solutions along the PDT
	for(int m=-5; m<6; m++) {
	  double cycle = m*wl;
	  double s[2]; //the two solutions in this cycle;
	  s[0] = xpos+cycle;
	  s[1] =  xpos2+cycle;
	  	

	  //now loop over the two solutions:
	  for(int j=0; j<2; j++) {
	
	    //require the solution is inside the PDT:
	    if(s[j]<posmax && s[j]>posmin) {
	
	      //calculate the residual to the segment position:
	      double res = abs(s[j] - segpos);
	
	      //keep the n best solutions
	      for(int k=0; k<nsol; k++) {
		if(res < abs(pbest[plane][k]-segpos)) {
		  double temp = pbest[plane][k];
		  pbest[plane][k]=s[j];
		  s[j] = temp;
		  res = abs(s[j] - segpos);
		}
	      }
	    } //inside the pdt
	  } //loop over 2 solutions per cycle
	
	} //loop over 11 cycles in one deck
	
       }//pdt pntr is valid
     } //loop over pdt links

   //now we have the two best solutions for each deck.
   //build a vector of all possible combinations of solutions in all 4 decks:

   vector<double> newvec, oldvec;
   vector<double> newvec2, oldvec2;
   vector<double> newvecsq, oldvecsq;
   vector<double> newvecsq2, oldvecsq2;
   ndeck=0;

   //reset the vectors
   newvec.push_back(0);
   newvec2.push_back(0);
   newvecsq.push_back(0);
   newvecsq2.push_back(0);
   oldvec.clear();
   oldvec2.clear();
   oldvecsq.clear();
   oldvecsq2.clear();


   //loop over the 4 decks:
   for(int i=0; i<4; i++) {

     //for each deck, we build a new vector of sums and squared sums.
     //copy the solutions so far into the old vector:
     oldvec = newvec;
     oldvec2 = newvec2;
     oldvecsq = newvecsq;
     oldvecsq2 = newvecsq2;
     newvec.clear();
     newvec2.clear();
     newvecsq.clear();
     newvecsq2.clear();

     //loop over solutions stored in this deck:
     int nadded=0;
     for(int j=0; j<nsol; j++) {
       if(pbest[i][j]>-998) {
	nadded++;
        // project pbest on the average plane, and call it qbest
        double offset = wbest[i] - planepos;
        double angle = pbest[i][j]/wbest[i];
        qbest[i][j] = pbest[i][j] - angle*offset;
	for (int k=0; k<oldvec.size(); k++) {
	  newvec.push_back( oldvec[k] + pbest[i][j] );
	  newvec2.push_back( oldvec2[k] + qbest[i][j] );
	  newvecsq.push_back( oldvecsq[k] + pbest[i][j]*pbest[i][j]);
	  newvecsq2.push_back( oldvecsq2[k] + qbest[i][j]*qbest[i][j]);
	}
       }	
     }//loop over n solutions

     //if we didn't add any solutions for this deck, copy the previous 
     //list of solutions into the new vector.
     if(nadded==0) {
      newvec = oldvec;
      newvec2 = oldvec2;
      newvecsq = oldvecsq;
      newvecsq2 = oldvecsq2;
     }
     else ndeck++;

   }//loop over decks


   if(ndeck==0) continue;


   rms = 999;
   newsegpos = -999;
   acceptpads = false;

   //if there is only one deck. take the first (closest) solution and 
   //set the RMS to 5 cm
   if(ndeck==1){
     float chi2 = (newvec2[0]-segpos)*(newvec2[0]-segpos)/(segerror*segerror);
     if (chi2<chi2max) {
       newsegpos = newvec2[0];
       rms=5;
       acceptpads = true;
     }
   }

   //for more than one deck, take the average and return the solution 
   //with the lowest RMS
   if(ndeck>1){
     for(int m=0; m<newvec.size(); m++) {
       double mean = newvec2[m]/ndeck;
       double sig2 = newvecsq2[m]/ndeck - 
                                  newvec2[m]*newvec2[m]/(ndeck*ndeck) ;
       if (sig2 < 0.00001) sig2 = 0.00001;
       double sig = sqrt(sig2);
       float chi2 = (mean-segpos)*(mean-segpos)/(segerror*segerror);

       if(sig < rms){
        if (chi2<chi2max) {
	  rms = sig;
	  newsegpos = mean;
          acceptpads = true;
        }
       }
     }
   } //ndeck>0

// update segment
   if (acceptpads && (ndeck>=_minDecks)) {
      double decks = ndeck;
      double pad_error = 5./sqrt(decks);
      double sum_wts = 1./(segerror*segerror) + 1./(pad_error*pad_error);
      double new_pos = (segpos/(segerror*segerror)+newsegpos/(pad_error*
                                                          pad_error))/sum_wts;
      double new_pos_error = sqrt(1./sum_wts);
      if (oct==0 || oct==3 || oct==4 || oct==7) {
        segiter->set_y(new_pos);
        segiter->set_y_error(new_pos_error);
      } else {
        segiter->set_x(new_pos);
        segiter->set_x_error(new_pos_error);
      }
      double xseg = segiter->x();
      double yseg = segiter->y();
      double dxseg = segiter->x_error();
      double dyseg = segiter->y_error();
      double phi = atan2(yseg,xseg);
      if (phi<0) phi = phi + kinem::TWOPI;
      segiter->set_phi(phi);
      double pex = dxseg*yseg;
      double pey = dyseg*xseg;
      double phi_err = sqrt(pex*pex + pey*pey)/(xseg*xseg + yseg*yseg);
      segiter->set_phi_error(phi_err);
   }
 }
 return;
}

} // namespace Muon