#include <iostream>
#include <sstream>
#include <unistd.h>
#include <fstream>
#include <cometEventLoop.hxx>
#include "ICOMETEvent.hxx"
#include "COMETGeomId.hxx"
#include <TROOT.h>
#include <TStyle.h>
#include <TFile.h>
#include <TTree.h>
#include <TMath.h>
#include "IReconFGDHacked.hxx"

/// Uses FGD isolated reconstruction to determine internal alignment parameters
/// Convention of layer numbers is 0, 2, 4, ...28 for FGD1X layers, 1, 3, 5, ... 29 for FGD1Y, 30, 32, ...42 for FGD2X, 31...43 for FGD2Y
class IRunFGDIntAlignLoop: public COMET::ICOMETEventLoopFunction 
{
protected:
  IReconFGDHacked *fgd_recon;
  
public:
  IRunFGDIntAlignLoop();
  virtual ~IRunFGDIntAlignLoop() {}
  
  bool operator () (COMET::ICOMETEvent& event);
  void BeginFile(COMET::IVInputFile* input);	
  bool SetOption(std::string option, std::string value="");
  void Initialize();
  void Finalize(COMET::ICOMETOutput* output);
  
private:	
  TFile *out;
  TTree *tree;
  
  const char* outputRootFileName;
  
  // Automatically acquire values from supplied geometry files with -G option
  double fgd1XInt;
  COMET::IGeometryId fgdLayer0;
  double fgd1YInt;
  COMET::IGeometryId fgdLayer1;
  COMET::IGeometryId fgdLayer2;
  double fgd2XInt;
  COMET::IGeometryId fgdLayer30;
  double fgd2YInt;
  COMET::IGeometryId fgdLayer31;
  COMET::IGeometryId fgdLayer32;
  
  double fgd1SepDist;
  double fgd2SepDist;
  
  double pos[44][10];
  double chg[44][10];
  double nodPos[44];
  double sumWgtPos[44];
  double sumChg[44];
  double wgtDiffSq[44];
  double wgtAvgPos[44];
  
  // Parameters to be saved	
  double res[44];
  int layerHitCnt[44];
  
  bool fgd1Xgood;
  bool fgd1Ygood;
  bool fgd2Xgood;
  bool fgd2Ygood;
  int evtno;
  int evtHitCnt;
  
  double fgd1TrackAngleX;
  double fgd1TrackAngleY;
  double fgd2TrackAngleX;
  double fgd2TrackAngleY;
  
  // Components of hits/nodes
  double nodX;
  double hitX;
  double nodY;
  double hitY;
  double nodZ;
  double hitZ;
  double charge;
  
  int hitItrPos;
  int layerNum;
  
  void FillArray(int start, int end, bool goodTrack)
  {
    layerNum = start;
    if (goodTrack)
      for (int i = start; i < end; i+=2)
  {
    wgtAvgPos[i] = sumWgtPos[i]/sumChg[i];
    res[layerNum] =  nodPos[i] - wgtAvgPos[i];
    //if(res[layerNum] > 10)
    //std::cout<<"Event "<<evtno<<", layer "<<layerNum<<" residual: "<<res[layerNum]<<std::endl;
    layerNum+=2;
  }
    else
      for (int i = start; i < end; i+=2)
  {
    res[layerNum] = 9999;
    layerNum+=2;
  }
  }
};	

/// Setup the parameters for rejection criteria.
void IRunFGDIntAlignLoop::Initialize()
{
  // Specify the output file name	
  outputRootFileName = COMET::IOARuntimeParameters::Get().GetParameterS("oaAlignTools.FGDIntAlignOutput").c_str();
  
  // Make/call output file and TTree inside.
  out = new TFile(outputRootFileName);
  if(out->IsOpen())
    {
      out->Close();
      out = new TFile(outputRootFileName, "update");
      tree = (TTree*)out->Get("fgdintaligntree");
      tree->SetBranchAddress("layer", &layerNum);
      tree->SetBranchAddress("event hit count", &evtHitCnt);
      tree->SetBranchAddress("residual", res);
      tree->SetBranchAddress("layer hit count", layerHitCnt);
      tree->SetBranchAddress("FGD1X track angle", &fgd1TrackAngleX);
      tree->SetBranchAddress("FGD1Y track angle", &fgd1TrackAngleY);
      tree->SetBranchAddress("FGD2X track angle", &fgd2TrackAngleX);
      tree->SetBranchAddress("FGD2Y track angle", &fgd2TrackAngleY);
    }
  else
    {
      out = new TFile(outputRootFileName, "recreate");	
      tree = new TTree("fgdintaligntree", "FGD internal alignment");
      tree->Branch("layer", &layerNum, "layerNum/I"); 
      tree->Branch("event hit count", &evtHitCnt, "evtHitCnt/I");
      tree->Branch("residual", res, "res[layerNum]/D");
      tree->Branch("layer hit count", layerHitCnt, "layerHitCnt[layerNum]/I");
      tree->Branch("FGD1X track angle", &fgd1TrackAngleX, "fgd1TrackAngleX/D");
      tree->Branch("FGD1Y track angle", &fgd1TrackAngleY, "fgd1TrackAngleY/D");
      tree->Branch("FGD2X track angle", &fgd2TrackAngleX, "fgd2TrackAngleX/D");
      tree->Branch("FGD2Y track angle", &fgd2TrackAngleY, "fgd2TrackAngleY/D");
    }
}
/// Write the TTree onto the file.	
void IRunFGDIntAlignLoop::Finalize(COMET::ICOMETOutput* output)
{
  delete fgd_recon;
  out->Write();
  out->Close();
}

IRunFGDIntAlignLoop::IRunFGDIntAlignLoop() 
{
  // Set the default output level
  //COMET::ICOMETLog::SetLogLevel("tpcAlign",COMET::ICOMETLog::QuietLevel);
}

// I have never used this... can I get rid of it?
bool IRunFGDIntAlignLoop::SetOption(std::string option, std::string value) 
{
  if (option == "v")
    {
      int value2 = atoi(value.c_str());
      if (value2 == 1)
  COMET::ICOMETLog::SetLogLevel("tpcAlign",COMET::ICOMETLog::LogLevel);
      else if (value2 == 2) 
  COMET::ICOMETLog::SetLogLevel("tpcAlign",COMET::ICOMETLog::InfoLevel);
      else if (value2 >= 3) 
  COMET::ICOMETLog::SetLogLevel("tpcAlign",COMET::ICOMETLog::VerboseLevel);
    }
  else return false;

return true;
}

/// Access geometry for module spacing and FGD intercepts
void IRunFGDIntAlignLoop::BeginFile(COMET::IVInputFile*)
{
  fgd_recon = new IReconFGDHacked();
  
}   

/// Main part of the program. Calls FGD Iso-recon and computes residuals
/// by comparing hit positions with node positions on the fitted tracks.
/// Charge on the hits is used as weight to perform weighted averaging of the 
/// hit position. Rejection criteria are all applied here.
bool IRunFGDIntAlignLoop::operator () (COMET::ICOMETEvent& event)
{
  
  COMET::IOADatabase::Get().Geometry(); 
  
  fgdLayer0 = COMET::GeomId::FGD::Layer(0, 0, 0);
  fgdLayer1 = COMET::GeomId::FGD::Layer(0, 0, 1);
  fgdLayer2 = COMET::GeomId::FGD::Layer(0, 1, 0);
  fgdLayer30 = COMET::GeomId::FGD::Layer(1, 0, 0);
  fgdLayer31 = COMET::GeomId::FGD::Layer(1, 0, 1);
  fgdLayer32 = COMET::GeomId::FGD::Layer(1, 1, 0);
  
  fgd1SepDist = fgdLayer2.GetPosition().Z()-fgdLayer0.GetPosition().Z();
  fgd2SepDist = fgdLayer32.GetPosition().Z()-fgdLayer30.GetPosition().Z();
  
  fgd1XInt = fgdLayer0.GetPosition().Z();
  fgd1YInt = fgdLayer1.GetPosition().Z();
  fgd2XInt = fgdLayer30.GetPosition().Z();
  fgd2YInt = fgdLayer31.GetPosition().Z();

  evtno = event.GetEventId();
  evtHitCnt = 0;
  fgd1Xgood = true;
  fgd1Ygood = true;
  fgd2Xgood = true;
  fgd2Ygood = true;
  layerNum = 0;
  fgd1TrackAngleX = 9999;
  fgd1TrackAngleY = 9999;
  fgd2TrackAngleX = 9999;
  fgd2TrackAngleY = 9999;
  if (evtno%100 == 0)
   std::cout<<"Processing event " << evtno <<std::endl;
  
  // ReconFGD's main algorithm without TPC matching
  COMET::IHandle<COMET::IHitSelection> fgd = event.GetHitSelection("fgd");
  if(fgd)
    {
      for (COMET::IHitSelection::const_iterator hit = fgd->begin(); hit != fgd->end(); ++hit)
  {
    //if ((*hit)->IsXHit())
    //  std::cout<<"This is X-hit!"<<" "<<(*hit)->GetPosition().Z()<<std::endl;
    //else
    // std::cout<<"This is Y-hit!"<<" "<<(*hit)->GetPosition().Z()<<std::endl;
    evtHitCnt++;
  }
    }
  
  IFGDTruthInfo::Get().SetEventInfo(event);
  COMET::IHandle<COMET::IAlgorithmResult> FgdResult = COMET::IHandle<COMET::IAlgorithmResult>();
  
  //FgdResult = event.GetFit("ReconFGD");
  
      FgdResult = fgd_recon->HackedProcess(fgd);
      // if(FgdResult)
      //event.AddFit(FgdResult, "ReconFGD");
      if (!FgdResult)
  return true;
  
  for (int i = 0; i < 44; i++)
    {
      sumWgtPos[i] = 0;
      sumChg[i] = 0;
      res[i] = 0;
      layerHitCnt[i] = 0;
      wgtDiffSq[i] = 0;
      for (int j = 0; j < 10; j++)
  {
    pos[i][j] = 0;
    chg[i][j] = 0;
  }	
    }
  // Get 2 sets of tracks from the FGD algorithm result.
  COMET::IHandle<COMET::IReconObjectContainer> xzReconTracks = (FgdResult)->GetResultsContainer("xzFgdReconTracks");
  COMET::IHandle<COMET::IReconObjectContainer> yzReconTracks = (FgdResult)->GetResultsContainer("yzFgdReconTracks");
  
  // Only let containers with 2 or less tracks - 1 for FGD1, 1 for FGD2 at maximum
  if (xzReconTracks && xzReconTracks->size() <= 2)	
    for(COMET::IReconObjectContainer::iterator tr = xzReconTracks->begin(); tr != xzReconTracks->end(); tr++)
      { 
  int xLayer;
  double increment;
  COMET::IHandle<COMET::IReconTrack> xzTrack = *tr;
  COMET::IHandle<COMET::IHitSelection> xzHits = xzTrack->GetHits();

  double trackPosX = xzTrack->GetPosition().X();
  double trackPosZ = xzTrack->GetPosition().Z();
  double trackSlope = xzTrack->GetDirection().X()/xzTrack->GetDirection().Z();
  if (trackPosZ > 1000)
    {
      fgd2TrackAngleX = atan(trackSlope)*180/TMath::Pi();
      xLayer = 7;
      layerNum = 30;
      hitZ = fgd2XInt;
      increment = fgd2SepDist;
    }
  else
    {
      fgd1TrackAngleX = atan(trackSlope)*180/TMath::Pi();
      xLayer = 15;
      layerNum = 0;
      hitZ = fgd1XInt;
      increment = fgd1SepDist;
    }
  for (int i = 0; i < xLayer; i++)
    {
      for(COMET::IHitSelection::iterator hit = xzHits->begin(); hit != xzHits->end(); hit++)
        {
    if (fabs((*hit)->GetPosition().Z()-hitZ) < 1)
      {
        pos[layerNum][layerHitCnt[layerNum]] = (*hit)->GetPosition().X();
        chg[layerNum][layerHitCnt[layerNum]] = (*hit)->GetCharge(); 
        
        layerHitCnt[layerNum]++;
        xzHits->RemoveHit(*hit);
        hit--;
      }
        }
      nodPos[layerNum] = trackPosX + trackSlope * (hitZ - trackPosZ);
      //std::cout<<layerNum<<" "<<nodPos[layerNum]<<std::endl;
      layerNum+=2;
      hitZ+=increment;
    }
      }
  
  // Only let containers with 2 or less tracks - 1 for FGD1, 1 for FGD2 at maximum
  if (yzReconTracks && yzReconTracks->size() <= 2)
    for(COMET::IReconObjectContainer::iterator tr = yzReconTracks->begin(); tr != yzReconTracks->end(); tr++)
      { 
  int yLayer;
  double increment;
  COMET::IHandle<COMET::IReconTrack> yzTrack = *tr;
  COMET::IHandle<COMET::IHitSelection> yzHits = yzTrack->GetHits();
  
  double trackPosY = yzTrack->GetPosition().Y();
  double trackPosZ = yzTrack->GetPosition().Z();
  double trackSlope = yzTrack->GetDirection().Y()/yzTrack->GetDirection().Z();
  if (trackPosZ > 1000)
    {
      fgd2TrackAngleY = atan(trackSlope)*180/TMath::Pi();
      yLayer = 7;
      layerNum = 31;
      hitZ = fgd2YInt;
      increment = fgd2SepDist;
    }
  else
    {
      fgd1TrackAngleY = atan(trackSlope)*180/TMath::Pi();
      yLayer = 15;
      layerNum = 1;
      hitZ = fgd1YInt;
      increment = fgd1SepDist;
    }
  for (int i = 0; i < yLayer; i++)
    {
      for(COMET::IHitSelection::iterator hit = yzHits->begin(); hit != yzHits->end(); hit++)
        {
    if (fabs((*hit)->GetPosition().Z()-hitZ) < 1)
      {
        pos[layerNum][layerHitCnt[layerNum]] = (*hit)->GetPosition().Y();
        chg[layerNum][layerHitCnt[layerNum]] = (*hit)->GetCharge(); 
        
        layerHitCnt[layerNum]++;
        yzHits->RemoveHit(*hit);
        hit--;
      }
        }
      nodPos[layerNum] = trackPosY + trackSlope * (hitZ - trackPosZ);
      layerNum+=2;
      hitZ+=increment;
    }
      }
  
  for (int i = 0; i < 44; i++)
    {
      double tempPos = 0;
      for (int j = 0; j < layerHitCnt[i]; j++)
  {
    if (j == 0 || fabs(pos[i][j] - tempPos) < 10.0)
      {
        tempPos = pos[i][j];
        sumWgtPos[i]+=pos[i][j]*chg[i][j];
        sumChg[i]+=chg[i][j];
      }
    else 
      {
        if (i < 30 && i%2 == 0)
    fgd1Xgood = false;
        else if (i < 30)
    fgd1Ygood = false;
        else if (i%2 == 0)
    fgd2Xgood = false;
        else
    fgd2Ygood = false;
      }
    
  }	
    } 

  FillArray(0, 30, fgd1Xgood);	
  FillArray(1, 30, fgd1Ygood);
  FillArray(30, 44, fgd2Xgood);	
  FillArray(31, 44, fgd2Ygood);
 
  tree->Fill();
  return true;
} 


int main(int argc, char *argv[])  
{
  IRunFGDIntAlignLoop userCode;
  cometEventLoop(argc,argv,userCode);
}