#include "ICOMETAlignmentLookup.hxx"
#include "TRotation.h"
#include "COMETGeomId.hxx"
#include "IOADatabase.hxx"
#include "ICOMETEvent.hxx"
#include "ICOMETContext.hxx"
#include "TGeoManager.h"
#include "TMath.h"
#include "IOARuntimeParameters.hxx"
#include <IAlignment_Fgd_Internal.hxx>
#include <IAlignment_P0D_Internal.hxx>
#include <IAlignment_COMET_Survey.hxx>
#include <IAlignment_TPC_MM.hxx>
#include "oaApplyAlign_version.h"

namespace {
   ICOMETAlignmentLookup invisibleAlignmentLookup;
}

ICOMETAlignmentLookup::ICOMETAlignmentLookup(){
    
  COMETInfo("Registering ICOMETAlignmentLookup alignment");
  COMET::IOADatabase::Get().RegisterAlignmentLookup(this);

  fApplyAlignment = 0;
  fApplyTrackerSurvey = 0;
  fApplyFGDIntAlign   = 0;
  fApplyP0DIntAlign   = 0;
  fApplyMMSurvey      = 0;
  fTestFlag = 0;

  // Set blank alignment ID to start.
  fAlignmentId = COMET::IAlignmentId();

  fRun = 0;

}


void ICOMETAlignmentLookup::Clear(){
  fAlignmentId = COMET::IAlignmentId();
  fCorrections.clear();
  fSHA1.Reset();
}



bool ICOMETAlignmentLookup::CheckAlignment(const COMET::ICOMETEvent* const event){
  
  // To start, just do a simple check; if the current run number
  // matchs the run number for this event, then we don't need to do 
  // any further checks of the alignment.
  if(fRun !=  event->GetContext().GetRun() && 
     event->GetContext().GetRun() != COMET::ICOMETContext::Invalid){    
    return true;
  }  

  return false;
  
}




//*****************************************************
COMET::IAlignmentId ICOMETAlignmentLookup::StartAlignment(const COMET::ICOMETEvent* const event) {
//*****************************************************

  // Set the flags for whether or not to apply any or all alignments
  fApplyAlignment = COMET::IOARuntimeParameters::Get().GetParameterI("oaApplyAlign.ApplyAlignment");
  
  fApplyTrackerSurvey = COMET::IOARuntimeParameters::Get().GetParameterI("oaApplyAlign.ApplyTrackerSurvey");
  fApplyFGDIntAlign   = COMET::IOARuntimeParameters::Get().GetParameterI("oaApplyAlign.ApplyFGDIntAlign");
  fApplyP0DIntAlign   = COMET::IOARuntimeParameters::Get().GetParameterI("oaApplyAlign.ApplyP0DIntAlign");
  fApplyMMSurvey      = COMET::IOARuntimeParameters::Get().GetParameterI("oaApplyAlign.ApplyMMSurvey");

  fTestFlag = COMET::IOARuntimeParameters::Get().GetParameterI("oaApplyAlign.TestFlag");

  // Save the current run number (if this is a valid context)
  if( event->GetContext().GetRun() != COMET::ICOMETContext::Invalid)
    fRun = event->GetContext().GetRun();

  // Check if alignment has already been applied to this event.
  // Should never happen.
  // Just pass back original alignment ID.
  COMET::IAlignmentId currentID = event->GetAlignmentId();
  if(currentID.Valid()){
    COMETError("Error: this geometry already has an alignment, "<<
               "but the method ICOMETAlignmentLookup::StartAlignment has been called");
    throw COMET::EGeometryAlreadyHadAlignment();
  }  
  
 
  // Don't apply any alignments if we dealing with simulated data.
  if( !(event->GetContext().IsDetector())) return fAlignmentId;

  // Don't store any default alignments if parameters flag is not set.
  // This will allow easier testing of new alignments.
  if(!fApplyAlignment) return fAlignmentId;

  // Clear out the old alignments, reset alignment ID.
  fCalls = 0;
  Clear(); 

  // Now store all the relevant alignments, for different methods.
  
  if(fApplyTrackerSurvey)
    ApplyTrackerSurveyAlign(event->GetContext());   

  if(fApplyMMSurvey)
    ApplyMMSurveyAlign(event->GetContext());

  if(fApplyFGDIntAlign)
    ApplyFGDInternalAlign(event->GetContext());

  if(fApplyP0DIntAlign)
    ApplyP0DInternalAlign(event->GetContext());



  // Add some final inputs to hash, based on the test flag.
  // Then calculate the hash and save in alignment ID.
  fSHA1.Input(fTestFlag);
  unsigned int messageDigest[5];
  if (fSHA1.Result(messageDigest)) {
    fAlignmentId = COMET::IAlignmentId(messageDigest);
    // Add document string to alignment ID.
    char document[200];
    sprintf(document,"align_v%ir%i_%i_%i_%i_%i_%i_%i",oaApplyAlign_MAJOR_VERSION, oaApplyAlign_MINOR_VERSION,
            fApplyAlignment,fApplyTrackerSurvey,
            fApplyMMSurvey,fApplyFGDIntAlign,fApplyP0DIntAlign,fTestFlag);
    fAlignmentId.SetDocumentString(document);
    COMETLog("Calculated alignment ID = " << fAlignmentId.AsString());
  }


  return fAlignmentId;

}

//*****************************************************
TGeoMatrix* ICOMETAlignmentLookup::Align(const COMET::ICOMETEvent* const event,
				       COMET::IGeometryId& geomId) {
//*****************************************************

  TGeoMatrix* result = NULL;
  if (fCalls<fCorrections.size()) {

    geomId = fCorrections[fCalls].first;
    result = fCorrections[fCalls].second;

  }

  ++fCalls;
  return result;

}

//*****************************************************
void  ICOMETAlignmentLookup::AddTranslation(const COMET::IGeometryId& geomId, double dx, double dy, double dz) {
//*****************************************************

  TGeoMatrix* trans = new TGeoTranslation(dx,dy,dz);
  fCorrections.push_back(std::pair<COMET::IGeometryId,TGeoMatrix*>(geomId,trans));

}

//*****************************************************
void  ICOMETAlignmentLookup::AddRotation(const COMET::IGeometryId& geomId, double a, double b, double c) {
//*****************************************************

// Euler angles in degrees
  TGeoMatrix* rot = new TGeoRotation("align",a,b,c);
  fCorrections.push_back(std::pair<COMET::IGeometryId,TGeoMatrix*>(geomId,rot));


}

//*****************************************************
void  ICOMETAlignmentLookup::AddRotation(const COMET::IGeometryId& geomId, TVector3 axis, double angle) {
//*****************************************************

  // angle is in milliradians.
  //  TGeoMatrix* rot = new TGeoRotation("align",angle,0,0);
  TGeoRotation* rot = new TGeoRotation("align");

  double elem[9];

  TRotation rot2;
  rot2.Rotate(angle/1000.,axis);
  elem[0] = rot2.XX();
  elem[1] = rot2.XY();
  elem[2] = rot2.XZ();
  elem[3] = rot2.YX();
  elem[4] = rot2.YY();
  elem[5] = rot2.YZ();
  elem[6] = rot2.ZX();
  elem[7] = rot2.ZY();
  elem[8] = rot2.ZZ();


  rot->SetMatrix(elem);

  fCorrections.push_back(std::pair<COMET::IGeometryId,TGeoMatrix*>(geomId,rot));


}

//*****************************************************
void  ICOMETAlignmentLookup::AddTranslationRotation(const COMET::IGeometryId& geomId, double dx, double dy, double dz,double a, double b, double c) {
//*****************************************************


// Euler angles in degrees
  TGeoTranslation trans(dx,dy,dz);
  TGeoRotation rot("align",a,b,c);
  TGeoMatrix* corr = new  TGeoCombiTrans(trans,rot);
  fCorrections.push_back(std::pair<COMET::IGeometryId,TGeoMatrix*>(geomId,corr));

}

//*****************************************************
void  ICOMETAlignmentLookup::AddTranslationRotation(const COMET::IGeometryId& geomId, double dx, double dy, double dz, TVector3 axis, double angle) {
//*****************************************************


// angle is in milliradians.
  TGeoTranslation trans(dx,dy,dz);

  TGeoRotation rot;

  double elem[9];

  TRotation rot2;
  rot2.Rotate(angle/1000.,axis);
  elem[0] = rot2.XX();
  elem[1] = rot2.XY();
  elem[2] = rot2.XZ();
  elem[3] = rot2.YX();
  elem[4] = rot2.YY();
  elem[5] = rot2.YZ();
  elem[6] = rot2.ZX();
  elem[7] = rot2.ZY();
  elem[8] = rot2.ZZ();


  rot.SetMatrix(elem);

  TGeoMatrix* corr = new  TGeoCombiTrans(trans,rot);
  fCorrections.push_back(std::pair<COMET::IGeometryId,TGeoMatrix*>(geomId,corr));

}
//*****************************************************
void  ICOMETAlignmentLookup::AddTranslationRotation(const COMET::IGeometryId& geomId, TVector3 trans, TVector3 axis, double angle) {
//*****************************************************


// angle is in milliradians.
  TGeoTranslation transVec(trans.X(),trans.Y(),trans.Z());

  TGeoRotation rot;

  double elem[9];

  TRotation rot2;
  rot2.Rotate(angle/1000.,axis);
  elem[0] = rot2.XX();
  elem[1] = rot2.XY();
  elem[2] = rot2.XZ();
  elem[3] = rot2.YX();
  elem[4] = rot2.YY();
  elem[5] = rot2.YZ();
  elem[6] = rot2.ZX();
  elem[7] = rot2.ZY();
  elem[8] = rot2.ZZ();


  rot.SetMatrix(elem);

  TGeoMatrix* corr = new  TGeoCombiTrans(transVec,rot);
  fCorrections.push_back(std::pair<COMET::IGeometryId,TGeoMatrix*>(geomId,corr));

}

/// The FGD internal alignment is defined by a set of constants derived 
/// from the calibration database.
void ICOMETAlignmentLookup::ApplyFGDInternalAlign(COMET::ICOMETContext context)
{
  COMETLog("Preparing FGD Internal Alignment");

  // Make the database query, using context to get the right constants.
  COMET::IResultSetHandle<COMET::IAlignment_Fgd_Internal> rs(context);


  Int_t numRows(rs.GetNumRows());
  // Loop over all layers we retrieve, saving the alignment for each.
  for (Int_t irow = 0; irow<numRows; ++irow){
    const COMET::IAlignment_Fgd_Internal* shift = rs.GetRow(irow);
    
    // Need to convert from um to mm for translation
    AddTranslation(shift->GetGeomId(), 
                   shift->GetDX()/1000.0, shift->GetDY()/1000.0, shift->GetDZ()/1000.0);
    
    // Add contents of this alignment to alignment hash.
    shift->AddToHash(fSHA1);
  }  

}	

/// The P0D internal alignment is defined by a set of constants derived 
/// from the calibration database.
void ICOMETAlignmentLookup::ApplyP0DInternalAlign(COMET::ICOMETContext context)
{
  COMETLog("Preparing P0D Internal Alignment");

  // Make the database query, using context to get the right constants.
  COMET::IResultSetHandle<COMET::IAlignment_P0D_Internal> rs(context);


  Int_t numRows(rs.GetNumRows());
  // Loop over all layers we retrieve, saving the alignment for each.
  for (Int_t irow = 0; irow<numRows; ++irow){
    const COMET::IAlignment_P0D_Internal* shift = rs.GetRow(irow);
    
    // Need to convert from um to mm for translation
    AddTranslation(shift->GetGeomId(), 
                   shift->GetDX()/1000.0, shift->GetDY()/1000.0, shift->GetDZ()/1000.0);
    
    // Add contents of this alignment to alignment hash.
    shift->AddToHash(fSHA1);
  }  

}	


#include <COMETGeomIdDef.hxx>
/// This method is copied out of oaGeomInfo.  Can't call the oaGeomInfo
/// version of this method, since we end up with nasty, ugly infinite recursion...
int BlahGeomIdToMM(const COMET::IGeometryId& id) {

  if(COMET::GeomId::TPC::IsMicroMega(id) || COMET::GeomId::TPC::IsPad(id))	
    return id.GetField(COMET::GeomId::Def::TPC::Pad::kMMegaMSB,
		       COMET::GeomId::Def::TPC::Pad::kMMegaLSB);
  
  else return -1;
}

/// This method is copied out of oaGeomInfo.  Can't call the oaGeomInfo
/// version of this method, since we end up with nasty, ugly infinite recursion...
int BlahGeomIdToHalf(const COMET::IGeometryId& id) {

  int detector = id.GetField(COMET::GeomId::Def::kDetectorIdMSB,
			     COMET::GeomId::Def::kDetectorIdLSB);
  if (detector != COMET::GeomId::Def::kTPC) return -1;
  /// This may be a pad.
  int seqId = id.GetField(COMET::GeomId::Def::TPC::kSeqIdMSB,
			  COMET::GeomId::Def::TPC::kSeqIdLSB);
  if (seqId != COMET::GeomId::Def::TPC::kPad) return -1;

  return id.GetField(COMET::GeomId::Def::TPC::Pad::kHalfMSB,
		     COMET::GeomId::Def::TPC::Pad::kHalfLSB);

}

/// Reads in TPC micromegas survey alignment constants and modifies geometry.
/// First column is the sub-detector geometry ID.
/// Second, third and fourth column are the translation constants (dx, dy, dz) in um.
/// Fifth column is the angle of rotation given in microradians.
void ICOMETAlignmentLookup::ApplyMMSurveyAlign(COMET::ICOMETContext context)
{

  COMETLog("Preparing MM Survey Alignment");

  COMET::IResultSetHandle<COMET::IAlignment_TPC_MM> rs(context);

  Int_t numRows(rs.GetNumRows());

  // Loop over all alignments in database.
  for (Int_t irow = 0; irow<numRows; ++irow){    
  
    const COMET::IAlignment_TPC_MM* shift = rs.GetRow(irow);

    // Need translation vector (convert from um to mm).
    TVector3 trans(shift->GetDX()/1000.0, shift->GetDY()/1000.0, shift->GetDZ()/1000.0);
    
    // Need theta in milliradians (convert from microradians)
    double thetaMRad = shift->GetAngle()/1000.0;

    // Need the endplate and MM #
    int mm = BlahGeomIdToMM(shift->GetGeomId());
    int endplate = BlahGeomIdToHalf(shift->GetGeomId());

    
    // Different coordinate system for every 6 MM's
    TVector3 mmAxis;
    if (endplate == 0 && mm >= 0 && mm < 6)
      {
        mmAxis.SetXYZ(0, 0, 1);
        AddTranslationRotation(shift->GetGeomId(), 
                               -trans.Z(), 
                               trans.Y(), 
                               trans.X(), 
                               mmAxis, 
                               thetaMRad);
      }
    else if (endplate == 0 && mm >= 6 && mm < 12)
      {
        mmAxis.SetXYZ(0, 0, 1);
        AddTranslationRotation(shift->GetGeomId(), 
                               trans.Z(), 
                               -trans.Y(), 
                               trans.X(), 
                               mmAxis, 
                               thetaMRad);	
      }	
    else if (endplate == 1 && mm >= 0 && mm < 6)
      {
        mmAxis.SetXYZ(0, 0, -1);
        AddTranslationRotation(shift->GetGeomId(), 
                               trans.Z(), 
                               trans.Y(),
                               -trans.X(), 
                               mmAxis, 
                               thetaMRad);
      }
    else
      {
        mmAxis.SetXYZ(0, 0, -1);
        AddTranslationRotation(shift->GetGeomId(), 
                               -trans.Z(), 
                               -trans.Y(), 
                               -trans.X(), 
                               mmAxis, 
                               thetaMRad);
      }	    


    // Add contents of this alignment to alignment hash.
    shift->AddToHash(fSHA1);

  }
  


}
/// This is Kendall's method to align TPC micromegas modules. 
/// This method requires a change in IOADatabase.cxx in oaEvent/src, which resolves the local-global coordinate system difference.
/*	
void ICOMETAlignmentLookup::ApplyMMSurveyAlignGlobal()
{
  int tpcNum;
  int rpNum;
  int mmNum;
  
  // No need to define rotation axis on case-by-case basis
  TVector3 mmAxis(1, 0, 0);
  
  alignLookupMM.open(inputMMSurveyAlignTxt);
  if (alignLookupMM.is_open())
    {
      std::cout<<"Micromegas survey alignment text detected: "<<inputMMSurveyAlignTxt<<std::endl;
      while (alignLookupMM >>  tpcNum >> rpNum >> mmNum >> transX >> transY >> transZ >> thetaMRad)
	{
	  
	  AddTranslationRotation(tpcMMId[(tpcNum-1)][rpNum][mmNum], transX, transY, transZ, mmAxis, thetaMRad);
	  std::cout<<"Alignment constants on TPC "<< tpcNum <<", RP "<<rpNum<<", MM " <<mmNum<<": ("<<
	    transX<<", "<<
	    transY<<", "<<
	    transZ<<", " <<
	    thetaMRad<<" mrad)"<<std::endl;
	  
	}
      alignLookupMM.close();
    }
  else
    std::cout<<"Micromegas survey alignment text " <<inputMMSurveyAlignTxt<<" not found, alignment not applied"<<std::endl;	
  
}	
*/	


/// Reads in COMET survey alignment constants from offline database.
///
/// First column is the sub-detector geometry ID
/// Second, third and fourth column are the translation constants (dx, dy, dz) in um
/// Fifth, sixth and seventh column are the normalized components of axis of rotation (ux, uy, uz) in um
/// Eighth column is the angle of rotation given in microradians.
/// 9th, 10th and 11th column is the offset between active and overall detector volumes in um.
void ICOMETAlignmentLookup::ApplyTrackerSurveyAlign(COMET::ICOMETContext context)
{
  COMETLog("Preparing Tracker-DsECAL Survey Alignment");

  // Query the database to get the TPC/FGD/DsECAL survey constants.
  COMET::IResultSetHandle<COMET::IAlignment_COMET_Survey> rs(context);

  Int_t numRows(rs.GetNumRows());

  for (Int_t irow = 0; irow<numRows; ++irow){
    //rs.GetRow(irow)->Print();
    const COMET::IAlignment_COMET_Survey* shift = rs.GetRow(irow);
    //AddTranslation(shift->GetGeomId(), shift->GetDX(), shift->GetDY(), shift->GetDZ());    
    
    // Need to convert distances from um (stored in database) to mm (expected by ROOT geometry)
    TVector3 trans(shift->GetDX()/1000.0, shift->GetDY()/1000.0, shift->GetDZ()/1000.0);
    //TVector3 trans;
    
    TVector3 rotAxis(shift->GetUX()/1000.0, shift->GetUY()/1000.0, shift->GetUZ()/1000.0);
    double mag = rotAxis.Mag();
    rotAxis = TVector3(rotAxis.X()/mag,
                       rotAxis.Y()/mag,
                       rotAxis.Z()/mag);
    
    // Need theta in radians and milliradians (convert from microradians)
    double thetaRad =  shift->GetAngle()/1.0e6;
    double thetaMRad = shift->GetAngle()/1000.0;

    TVector3 activeVolumeOffset(shift->GetAVX()/1000.0, shift->GetAVY()/1000.0, shift->GetAVZ()/1000.0);

    // Dean's rotation matrix.
    // This will, I think, only work for shifts in the Y-axis between active and overall volumes.  Warn if
    // otherwise
    if(activeVolumeOffset.X() != 0.0 || activeVolumeOffset.Z() != 0.0)
      COMETLog("WARNING: the hardcoded transformations that account for the difference between active\n" <<
               "and overall volume only works for shifts in Y.\nYou have added extra shifts in X (" << 
               activeVolumeOffset.X() << ") and Z ("  << activeVolumeOffset.Z() << ").\n" <<
               "This will not be correctly handled.");

    double trackerRotMatrix[9];
    trackerRotMatrix[0] = 2*pow(sin(thetaRad/2),2)*pow(rotAxis.X(),2) + cos(thetaRad);
    trackerRotMatrix[1] = 2*pow(sin(thetaRad/2),2)*rotAxis.X()*rotAxis.Y() - sin(thetaRad)*rotAxis.Z();
    trackerRotMatrix[2] = 2*pow(sin(thetaRad/2),2)*rotAxis.X()*rotAxis.Z() + sin(thetaRad)*rotAxis.Y();
    trackerRotMatrix[3] = 2*pow(sin(thetaRad/2),2)*rotAxis.X()*rotAxis.Y() + sin(thetaRad)*rotAxis.Z();
    trackerRotMatrix[4] = 2*pow(sin(thetaRad/2),2)*pow(rotAxis.Y(),2) + cos(thetaRad);
    trackerRotMatrix[5] = 2*pow(sin(thetaRad/2),2)*rotAxis.Y()*rotAxis.Z() - sin(thetaRad)*rotAxis.X();
    trackerRotMatrix[6] = 2*pow(sin(thetaRad/2),2)*rotAxis.X()*rotAxis.Z() - sin(thetaRad)*rotAxis.Y();
    trackerRotMatrix[7] = 2*pow(sin(thetaRad/2),2)*rotAxis.Y()*rotAxis.Z() + sin(thetaRad)*rotAxis.X();
    trackerRotMatrix[8] = 2*pow(sin(thetaRad/2),2)*pow(rotAxis.Z(),2) + cos(thetaRad);
    
    TVector3 offset(trackerRotMatrix[1]*-activeVolumeOffset.Y(), 
                    (trackerRotMatrix[4]-1)*-activeVolumeOffset.Y(), 
                    trackerRotMatrix[7]*-activeVolumeOffset.Y());

    AddTranslationRotation(shift->GetGeomId() ,
                         offset + trans, 
                         rotAxis, 
                         thetaMRad);

    // Add contents of this alignment to alignment hash.
    shift->AddToHash(fSHA1);

  }

}