//COMET Files
#include <cometEventLoop.hxx>
#include <IChargeClusters2D.hxx>
#include <IP0D3DTrackMatching.hxx>
#include <IGeomInfo.hxx>
#include <IComboHit.hxx>
#include <IFieldManager.hxx>

//root files
#include <TFile.h>
#include <TTree.h>

//C++ files
#include <cstdio>

/// Uses P0D Isolated Reconstruction to determine P0D internal alignment coefficients.
class IRunP0DIntAlignLoop: public COMET::ICOMETEventLoopFunction {
public:
  IRunP0DIntAlignLoop() {}
  virtual ~IRunP0DIntAlignLoop() {}

  bool operator () (COMET::ICOMETEvent&);
  void Initialize();
  void Finalize(COMET::ICOMETOutput*);

  //Create struct to temporarily save hit information
  struct HitInfo {
    double pos;
    double zpos;
    double charge;
    int size;
  } Hit;

  //Functions for P0D internal alignment
  std::vector<std::vector<double> > GetDeltaMatrix();
  bool IsThroughTrack(COMET::IHandle<COMET::IReconTrack>);
  bool IsClose(COMET::IHandle<COMET::IHitSelection>& , COMET::IHandle<COMET::IHit>&);
  COMET::IHandle<COMET::IHitSelection> AddNearbyHits(COMET::IHandle<COMET::IHitSelection>&, COMET::IHandle<COMET::IHitSelection>&);
  COMET::IHandle<COMET::IHitSelection> SaveAdjacentHits(COMET::IHandle<COMET::IHitSelection>&);
  bool IsEveryLayer(COMET::IHandle<COMET::IHitSelection>&);
  std::pair<double,double> LinearFit(std::map<int,HitInfo>);
  double GetChiSqr(std::map<int,HitInfo>, std::pair<double,double>, bool);
  int FindNonet(HitInfo, std::pair<double,double>, bool);

private:
  //Number of tracks used
  double fCountTracks;

  //File to be saved
  TFile* fRootFile;

  //Trees to be saved
  TTree* fRootTreeX;
  TTree* fRootTreeY;

  //Variables to be saved
  double fP0Dule;
  double fCharge;
  double fRes;
  double fAngle;
  double fChi;
  double fNonet;
  int fSize;

  //Matrix of previous alignment parameters
  std::vector<std::vector<double> > fDeltaMatrix;
};

///Function that retrieves the previous alignment constants during iteration.
///If no file called P0DInt.txt is present in working directory it returns a matrix of zeros.
std::vector<std::vector<double> > IRunP0DIntAlignLoop::GetDeltaMatrix() {
  std::vector<std::vector<double> > output;
  output.erase(output.begin(), output.end());
  std::vector<double> column;
  column.erase(column.begin(), column.end());
  column.assign(2,0);
  output.assign(40, column);
  
  int layerNum=0; float x=0; float ex=0; float y=0; float ey=0;
  std::FILE * P0DIntInput = std::fopen ("P0DInt.txt","r");
  if (P0DIntInput == NULL) {
    COMETLog("No Alignment Parameters to Apply: Returning Zero Array");
    return output;
  }
  else {
    for (int i=0; i<40; i++) {
      std::fscanf (P0DIntInput, "%i     %f      %f      %f      %f", &layerNum, &x, &ex, &y, &ey);
      output[i][0] = x;
      output[i][1] = y;
    }
    std::fclose (P0DIntInput);
    COMETLog("Read in previous parameters");
  }
  return output;
}


///Function that returns true if the track passes through the first and last layers of the P0D.
bool IRunP0DIntAlignLoop::IsThroughTrack(COMET::IHandle<COMET::IReconTrack> p0dTrack) {
  //Check if there is a Hit in the First and Last Layer
  COMET::IHitSelection::iterator firstHit = (p0dTrack->GetNodes()).front()->GetObject()->GetHits()->begin();
  int firstP0Dule = -1;
  if((*firstHit)->IsYHit()) firstP0Dule = COMET::IGeomInfo::Get().P0D().ActiveYPlane((*firstHit)->GetPosition().Z());
  if((*firstHit)->IsXHit()) firstP0Dule = COMET::IGeomInfo::Get().P0D().ActiveXPlane((*firstHit)->GetPosition().Z());
  COMET::IHitSelection::iterator lastHit = (p0dTrack->GetNodes()).back()->GetObject()->GetHits()->begin();
  int lastP0Dule = -1;
  if((*lastHit)->IsYHit()) lastP0Dule = COMET::IGeomInfo::Get().P0D().ActiveYPlane((*lastHit)->GetPosition().Z());
  if((*lastHit)->IsXHit()) lastP0Dule = COMET::IGeomInfo::Get().P0D().ActiveXPlane((*lastHit)->GetPosition().Z());
  
  //Length cut
  if (firstP0Dule != 0 || lastP0Dule != 39) return false;
  return true;
}

///Function returns true if an unused hit is within (40mm, 40mm, 1cm) to any other hit in a set of hits.
bool IRunP0DIntAlignLoop::IsClose(COMET::IHandle<COMET::IHitSelection>& trackHits, COMET::IHandle<COMET::IHit>& comparisonHit) {
  double testZ = comparisonHit->GetPosition().Z();
  for (COMET::IHitSelection::iterator hit = trackHits->begin(); hit != trackHits->end(); ++hit) {
    double tempZ = (*hit)->GetPosition().Z();
    if (tempZ < testZ+1*unit::cm && tempZ > testZ-1*unit::cm) {
      if (comparisonHit->IsXHit()) {
  double testX = comparisonHit->GetPosition().X();
  if ((*hit)->IsXHit()) {
    double tempX = (*hit)->GetPosition().X();
    if (tempX < testX+40*unit::mm && tempX > testX-40*unit::mm) return true;
  }
      }
      if (comparisonHit->IsYHit()) {
  double testY = comparisonHit->GetPosition().Y();
  if ((*hit)->IsYHit()) {
    double tempY = (*hit)->GetPosition().Y();
    if (tempY < testY+40*unit::mm && tempY > testY-40*unit::mm) return true;
  }
      }
    }
  }
  return false;
}

///Function adds nearby unused hits to track hits.
COMET::IHandle<COMET::IHitSelection> IRunP0DIntAlignLoop::AddNearbyHits(COMET::IHandle<COMET::IHitSelection>& objectHits, COMET::IHandle<COMET::IHitSelection>& unusedHits) {
  if (unusedHits->size() != 0) {
    for (COMET::IHitSelection::iterator hits = unusedHits->begin(); hits != unusedHits->end(); ++hits) {
      if (IsClose(objectHits,(*hits))) objectHits->AddHit((*hits));
    }
  }
  return objectHits;
}

///Saves hit selections with only singlet and doublet hits.
COMET::IHandle<COMET::IHitSelection> IRunP0DIntAlignLoop::SaveAdjacentHits(COMET::IHandle<COMET::IHitSelection>& inputHits) {  
  COMET::IHandle<COMET::IHitSelection> outputHits(new COMET::IHitSelection("outputHits"));
  for (COMET::IHitSelection::iterator hit = inputHits->begin(); hit != inputHits->end(); ++hit) {
    COMET::IHandle<COMET::IComboHit> chit = (*hit);
    if(!chit) COMETWarn("ComboHit Handle not Valid!!!");
    if (chit->GetHits().size() == 1) outputHits->AddHit(chit);
    if (chit->GetHits().size() == 2) {
      int barNumber0 = COMET::IGeomInfo::P0D().GetBar((*chit->GetHits().front()).GetGeomId()).GetNumber();
      int barNumber1 = COMET::IGeomInfo::P0D().GetBar((*chit->GetHits().back()).GetGeomId()).GetNumber();
      if (abs(barNumber0-barNumber1) == 1) outputHits->AddHit(chit);
    }
  }      
  return outputHits;
}

///Function returns true if there is at most one clustered hit in every layer.
bool IRunP0DIntAlignLoop::IsEveryLayer(COMET::IHandle<COMET::IHitSelection>& trackHits) {
  std::map<int,int> p0duleMap;
  for (int i = 0; i<40; i++) p0duleMap[i]=0;
  for (COMET::IHitSelection::iterator hit = trackHits->begin(); hit != trackHits->end(); hit++) {
    double z = (*hit)->GetPosition().Z();
    int layer;
    if ((*hit)->IsYHit()) layer = COMET::IGeomInfo::Get().P0D().ActiveYPlane(z);
    if ((*hit)->IsXHit()) layer = COMET::IGeomInfo::Get().P0D().ActiveXPlane(z);
    p0duleMap[layer]++;
  }
  for (std::map<int,int>::iterator it = p0duleMap.begin(); it != p0duleMap.end(); it++) {
    if (it->second > 1) return false;
  }
  return true;
}

///Function fits a straight line to a collection of hits.
std::pair<double,double> IRunP0DIntAlignLoop::LinearFit(std::map<int,HitInfo> inputMap) {
  //Initializes values to be used in the fitting function.;
  std::pair<double,double> output;
  double sumX = 0;
  double sumY = 0;
  double sumXY = 0;
  double sumXX = 0;
  int size = 0;
  for (std::map<int,HitInfo>::iterator it = inputMap.begin(); it != inputMap.end(); it++) {
    double y = (*it).second.pos;
    double x = (*it).second.zpos;
    sumX += x;
    sumY += y;
    sumXY += x * y;
    sumXX += x * x;
    size += 1;
  }
  //from Bevington page 105 Third Edition
  output.first = (size * sumXY - sumX * sumY)/(size * sumXX - sumX * sumX);
  output.second = (sumXX * sumY - sumX * sumXY)/(size * sumXX - sumX * sumX);
  return output;
}

///Function to find the Chi Square of the fit line.
double IRunP0DIntAlignLoop::GetChiSqr(std::map<int,HitInfo> hitMap, std::pair<double,double> fitLine, bool isXMap) {
  double output;
  double sumSqr = 0;
  double dof = -2;
  for (std::map<int,HitInfo>::iterator it = hitMap.begin(); it != hitMap.end(); it++) {
    sumSqr += pow((*it).second.pos - fitLine.first * (*it).second.zpos - fitLine.second + fDeltaMatrix[(*it).first][isXMap],2);
    dof += 1;
  }
  output = sumSqr/dof;
  return output;
}

///Finds the nonet that the hit belongs to.
///This will be used for angular alignment if needed.
int IRunP0DIntAlignLoop::FindNonet(HitInfo Hit, std::pair<double,double> fitLine, bool XorY) {
  double x,y;
  if (XorY) {
    x = Hit.pos;
    y = fitLine.first * Hit.zpos + fitLine.second;
  }
  else {
    x = fitLine.first * Hit.zpos + fitLine.second;
    y = Hit.pos;
  }
  double xMin = COMET::IGeomInfo::P0D().ActiveMin().X();
  double xMax = COMET::IGeomInfo::P0D().ActiveMax().X();
  double yMin = COMET::IGeomInfo::P0D().ActiveMin().Y();
  double yMax = COMET::IGeomInfo::P0D().ActiveMax().Y();
  int result = -1;
  if (y > (2*yMax+yMin)/3) {
    if (x < (xMax+2*xMin)/3) result = 0;
    if (x > (xMax+2*xMin)/3 && x < (2*xMax+xMin)/3) result = 1;
    if (x > (2*xMax+xMin)/3) result = 2;
  }
  if (y < (2*yMax+yMin)/3 && y > (yMax+2*yMin)/3) {
    if (x < (xMax+2*xMin)/3) result = 3;
    if (x > (xMax+2*xMin)/3 && x < (2*xMax+xMin)/3) result = 4;
    if (x > (2*xMax+xMin)/3) result = 5;
  }
  if (y < (yMax+2*yMin)/3) {
    if (x < (xMax+2*xMin)/3) result = 6;
    if (x > (xMax+2*xMin)/3 && x < (2*xMax+xMin)/3) result = 7;
    if (x > (2*xMax+xMin)/3) result = 8;
  }
  return result;
}

///Setup files and parameters.
void IRunP0DIntAlignLoop::Initialize() {
  fRootFile = new TFile("P0DAlignParam.root","recreate");
  fRootTreeX = new TTree("ResidualX","Residual X direction");
  fRootTreeX->Branch("p0dule",&fP0Dule,"fP0Dule/D");
  fRootTreeX->Branch("residual",&fRes,"fRes/D");
  fRootTreeX->Branch("charge",&fCharge,"fCharge/D");
  fRootTreeX->Branch("angle",&fAngle,"fAngle/D");
  fRootTreeX->Branch("chi",&fChi,"fChi/D");
  fRootTreeX->Branch("nonet",&fNonet,"fNonet/D");
  fRootTreeX->Branch("size",&fSize,"fSize/I");
  fRootTreeY = new TTree("ResidualY","Residual Y direction");
  fRootTreeY->Branch("p0dule",&fP0Dule,"fP0Dule/D");
  fRootTreeY->Branch("residual",&fRes,"fRes/D");
  fRootTreeY->Branch("charge",&fCharge,"fCharge/D");
  fRootTreeY->Branch("angle",&fAngle,"fAngle/D");
  fRootTreeY->Branch("chi",&fChi,"fChi/D");
  fRootTreeY->Branch("nonet",&fNonet,"fNonet/D");
  fRootTreeY->Branch("size",&fSize,"fSize/I");
  fCountTracks = 0;
  fDeltaMatrix = GetDeltaMatrix();
}

/// Main part of program. Calls for ReconP0D then extracts tracks that cross
/// the entire P0D and finds residuals to the linear fit of the tracks.
bool IRunP0DIntAlignLoop::operator() (COMET::ICOMETEvent& event) {
  //opens geometry file associated with event file
  COMET::IOADatabase::Get().Geometry();
  
  // Check if magnetic field exists
  double B =  COMET::IFieldManager::GetFieldValue( TVector3(0,0,0) ).Mag()/unit::tesla;
  if (B != 0) {
    COMETError("Non-Zero Magnetic Field in Event!");
    return false;
  }

  // Get the 3D P0DRecon object
  COMET::IHandle<COMET::IReconObjectContainer> P0DTracks;
  COMET::IHandle<COMET::IAlgorithmResult> P0DCycle;
  COMET::IHandle<COMET::IReconBase> p0dTrack;
  bool thru = false;
  for (int cycleCount = 0; cycleCount < 23; cycleCount++) {
    //For Production 5 ReconP0D structure.
    P0DCycle = event.GetFit(TString::Format("ReconP0D/p0dCycle%d/IP0DTrackRecon/IP0D3DTrackMatching", cycleCount));
    //For pre-Production 5 ReconP0D structure.
    if (!P0DCycle)
      P0DCycle = event.GetFit(TString::Format("ReconP0D/p0dCycle%d/IP0D3DTrackMatching", cycleCount));
    if (!P0DCycle) continue;
    COMET::IHandle<COMET::IReconObjectContainer> P0DFinal = (P0DCycle)->GetResultsContainer("final");
    if (!P0DFinal || P0DFinal->size() != 1) continue;
    P0DTracks = P0DFinal;
    
    // Get the Track from the Recon object
    p0dTrack = P0DTracks->front();
    thru = IsThroughTrack(p0dTrack);
    if (!thru) continue;
    break;
  }
  if (!thru) return false;
 
  //unused hits in the track reconstruction
  COMET::IHandle<COMET::IHitSelection> unusedHits = P0DCycle->GetHitSelection("unused");
  bool unusedExists = true;
  if (!unusedHits) unusedExists = false;
  
  //Definition of Cluster
  COMET::IChargeClusters2D p0dCluster(30*unit::mm,30*unit::mm,10*unit::mm);
  
  bool hasGoodTrack = false;
  
  //Get the Hits from the Reconstructed Track
  COMET::IHandle<COMET::IHitSelection> trackHits = p0dTrack->GetHits();
  COMET::IHandle<COMET::IHitSelection> totalHits;
  
  //Adds missing hits if any from the Track
  if (unusedExists) totalHits = AddNearbyHits(trackHits,unusedHits);
  else totalHits = trackHits;
  
  //cluster hit selection
  COMET::IHandle<COMET::IHitSelection>  clusteredHits(p0dCluster(*totalHits));
  
  //Saves singlets and doublets in adjacent bars
  COMET::IHandle<COMET::IHitSelection> finalHits = SaveAdjacentHits(clusteredHits);
  
  //create 2D hit selections
  COMET::IHandle<COMET::IHitSelection> xHits(new COMET::IHitSelection("xHits"));
  COMET::IHandle<COMET::IHitSelection> yHits(new COMET::IHitSelection("yHits"));
  
  //separate 3D track into 2D RawHits
  for (COMET::IHitSelection::iterator hits = finalHits->begin(); hits != finalHits->end(); ++hits) {
    if ((*hits)->IsXHit()) xHits->push_back((*hits));
    if ((*hits)->IsYHit()) yHits->push_back((*hits));
  }    
  
  //Check if the Track has a most one clustered hit in every bar
  if (!IsEveryLayer(xHits)) return false;
  if (!IsEveryLayer(yHits)) return false;
  
  hasGoodTrack = true;
  fCountTracks += 1;
  
  //create maps of layers to hit information
  std::map<int,HitInfo> XMap;
  std::map<int,HitInfo> YMap;
  
  //add Information to Maps
  int layer;
  for (COMET::IHitSelection::iterator hits = xHits->begin(); hits != xHits->end(); ++hits) {
    COMET::IHandle<COMET::IComboHit> chit = (*hits);
    if(!chit) COMETWarn("ComboHit Handle not Valid!!!");
    Hit.size = chit->GetHits().size();
    Hit.pos = (*hits)->GetPosition().X();
    Hit.zpos = (*hits)->GetPosition().Z();
    Hit.charge = (*hits)->GetCharge();
    layer = COMET::IGeomInfo::Get().P0D().ActiveXPlane(Hit.zpos);
    XMap[layer] = Hit;
  }
  
  for (COMET::IHitSelection::iterator hits = yHits->begin(); hits != yHits->end(); ++hits) {
    COMET::IHandle<COMET::IComboHit> chit = (*hits);
    if(!chit) COMETWarn("ComboHit Handle not Valid!!!");
    Hit.size = chit->GetHits().size();
    Hit.pos = (*hits)->GetPosition().Y();
    Hit.zpos = (*hits)->GetPosition().Z();
    Hit.charge = (*hits)->GetCharge();
    layer = COMET::IGeomInfo::Get().P0D().ActiveYPlane(Hit.zpos);
    YMap[layer] = Hit;
  }
  
  //Linear Fits to P0D
  std::pair<double,double> fullP0DX = LinearFit(XMap);
  std::pair<double,double> fullP0DY = LinearFit(YMap);
  
  //Get Chi Square.
  double xChi = GetChiSqr(XMap,fullP0DX,true);
  double yChi = GetChiSqr(YMap,fullP0DY,false);
  
  //Fill P0D Information
  for (std::map<int,HitInfo>::iterator it = XMap.begin(); it != XMap.end(); it++) {
    //Fills the Internal Tree
    fSize = (*it).second.size;
    fAngle = atan(fullP0DX.first);
    fRes = (*it).second.pos - fullP0DX.first * (*it).second.zpos - fullP0DX.second + fDeltaMatrix[(*it).first][0];
    fCharge = (*it).second.charge;
    fP0Dule = (*it).first;
    fChi = xChi;
    fNonet = FindNonet((*it).second, fullP0DY, true);
    fRootTreeX->Fill();
  }
  for (std::map<int,HitInfo>::iterator it = YMap.begin(); it != YMap.end(); it++) {
    //Fills the Internal Tree
    fSize = (*it).second.size;
    fAngle = atan(fullP0DY.first);
    fRes = (*it).second.pos - fullP0DY.first * (*it).second.zpos - fullP0DY.second + fDeltaMatrix[(*it).first][1];
    fCharge = (*it).second.charge;
    fP0Dule = (*it).first;
    fChi = yChi;
    fNonet = FindNonet((*it).second, fullP0DX, false);
    fRootTreeY->Fill();
  }
  return hasGoodTrack;
}

/// Writes the TTrees to file
void IRunP0DIntAlignLoop::Finalize(COMET::ICOMETOutput* output) {
  
  COMETLog("The number of Tracks used for Tree Building is " << fCountTracks);
  
  fRootFile->Write();
  fRootFile->Close();
}

int main(int argc, char **argv) {
  IRunP0DIntAlignLoop userCode;
  cometEventLoop(argc,argv,userCode);
}