#include <iostream>
#include <cmath>
#include <algorithm>

#include <TGeoManager.h>
#include <TGeoNode.h>
#include <TVector3.h>
#include <TRandom.h>

#include <HEPUnits.hxx>

#include <ICOMETLog.hxx>
#include <IGeomIdManager.hxx>
#include "IStrawTrkGeom.hxx"
#include "IGeomInfo.hxx"

#include <IOADatabase.hxx>
#include <COMETGeomId.hxx>
#include <COMETGeomIdDef.hxx>
#include <EoaGeomInfo.hxx>

#include "IDetectorSolenoidGeom.hxx"

COMET::IStrawTrkGeom::IStrawTrkGeom() :
  IGeomBase(),
  fStrawTrkWireManager(NULL),
  fStrawDiameter(-1),
  fStrawWireDiameter(-1),
  fNumberOfStations(0),
  fNumberOfManifolds(0),
  fNumberOfStraws(0)
{
  fNodeIdMap.resize(0);
  fNodeId.resize(0);
}

COMET::IStrawTrkGeom::~IStrawTrkGeom() { Clear(); }


void COMET::IStrawTrkGeom::Fill()
{
  fStrawTrkWireManager = new COMET::IStrawTrkWireManager();
  IStrawTrkGeomVisitor visitor(this);

  visitor.VisitGeometry();

  // Dupulicated numbers 
  fNumberOfStraws     = fNodeId.size()/fNumberOfManifolds;
  fNumberOfManifolds /= fNumberOfStations;
    
  BuildNodeIdMap();
  fStrawTrkWireManager->Init(fNodeIdMap);
  
  COMETNamedInfo("IStrawTrkGeom", "// --------------------------------------------------------------------- ");
  COMETNamedInfo("IStrawTrkGeom", "// StrawTrk geometry ");
  COMETNamedInfo("IStrawTrkGeom", "// --------------------------------------------------------------------- ");
  COMETNamedInfo("IStrawTrkGeom", "//  Number of Stations  = " << GetNumberOfStations());
  COMETNamedInfo("IStrawTrkGeom", "//  Number of Manifolds = " << GetNumberOfManifolds());
  COMETNamedInfo("IStrawTrkGeom", "//  Number of Straws    = " << GetNumberOfStraws());
  COMETNamedInfo("IStrawTrkGeom", "//  The position measurement direction of the front layer is "
                 << (COMET::GeomId::StrawTrk::kStrawTrkDirectionX ==
                     fStrawTrkWireManager->GetManifoldIdOfFrontLayer()? "X":"Y"));
  COMETNamedInfo("IStrawTrkGeom", "//  Global layer Z position in DS coordinates as follows: ");
  for(Int_t glayer=0; glayer<GetNumberOfGlobalLayers(); glayer++){
    COMETNamedInfo("IStrawTrkGeom", "//    [" << glayer << "] "
                   << GetGlobalLayerPositionZInDSFromGlobalLayerId(glayer) << " mm.");
  }
  COMETNamedInfo("IStrawTrkGeom", "// --------------------------------------------------------------------- ");  
}

void COMET::IStrawTrkGeom::ls(const char *opt)
{
  std::cout << "// --------------------------------------------------------------------- " << std::endl
            << "// StrawTrk geometry " << std::endl
            << "// --------------------------------------------------------------------- " << std::endl
            << "//  Number of Stations  = " << GetNumberOfStations() << std::endl
            << "//  Number of Manifolds = " << GetNumberOfManifolds() << std::endl
            << "//  Number of Straws    = " << GetNumberOfStraws() << std::endl
            << "//  The position measurement direction of the front layer is "
            << (COMET::GeomId::StrawTrk::kStrawTrkDirectionX ==
                fStrawTrkWireManager->GetManifoldIdOfFrontLayer()? "X":"Y") << std::endl
            << "//  Global layer Z position in DS coordinates as follows: " << std::endl;
  for(Int_t glayer=0; glayer<GetNumberOfGlobalLayers(); glayer++){
    std::cout << "//    [" << glayer << "] "
              << GetGlobalLayerPositionZInDSFromGlobalLayerId(glayer) << " mm." << std::endl;
  }
  std::cout << "// --------------------------------------------------------------------- " << std::endl;
}

void COMET::IStrawTrkGeom::Clear(void)
{
  delete fStrawTrkWireManager;
  fStrawTrkWireManager = NULL;
}

void COMET::IStrawTrkGeom::BuildNodeIdMap()
{
  COMETNamedVerbose("IStrawTrkGeom", "COMET::IStrawTrkGeom::BuildNodeIdMap");

  std::vector<int> tmp1;
  std::vector< std::vector<int> > tmp2;
  tmp1.resize(0); tmp2.resize(0);
  for (unsigned int ii=0; ii<fNodeId.size(); ii++) {

    int node_id = fNodeId.at(ii);
    tmp1.push_back(node_id);

    if(tmp1.size()%fNumberOfStraws == 0){
      tmp2.push_back(tmp1);
      tmp1.resize(0);

      if(tmp2.size()%fNumberOfManifolds==0){
        fNodeIdMap.push_back(tmp2);
        tmp2.resize(0);
      }
    }
  }
}

// -------------------------------------------------------------------------------------
bool COMET::IStrawTrkGeomVisitor::VisitNode(int depth, const std::string& name, const TGeoNode* node)
{
  bool ret = true;
  
  /// If the node isn't in the Straw Tracker, skip this node.
  gGeoManager->cd(name.c_str());

#define ISTGIF_CONTAINS(str) ( name.find(str) != (std::string::npos) )
  if(ISTGIF_CONTAINS("StrawWire"     )){
    ret = false;
    // Get shape info.
    if(fStrawTrkGeom->fStrawWireDiameter < 0){
      TVector3 nodeSize = COMET::IOADatabase::Get().GeomId().NodeSize(node,COMET::IGeomIdManager::kTube);
      fStrawTrkGeom->fStrawWireDiameter = 2*nodeSize.X();
    }
  }
  else if(ISTGIF_CONTAINS("StrawGas"      ))
    fStrawTrkGeom->fNodeId.push_back(gGeoManager->GetCurrentNodeId());
  else if(ISTGIF_CONTAINS("StrawTube"     )){
    // Get shape info
    if(fStrawTrkGeom->fStrawDiameter < 0){
      TVector3 nodeSize = COMET::IOADatabase::Get().GeomId().NodeSize(node,COMET::IGeomIdManager::kTube);
      fStrawTrkGeom->fStrawDiameter = 2*nodeSize.X();
    }
  }
  else if(ISTGIF_CONTAINS("Layer")){
  }
  else if(ISTGIF_CONTAINS("ManifoldBarrel") ||
          ISTGIF_CONTAINS("ManifoldCap"   ) ||    
          ISTGIF_CONTAINS("ROESTI"        )) ret = false;
  else if(ISTGIF_CONTAINS("Manifold"      )) {
    fStrawTrkGeom->fNumberOfManifolds++;
  }
  else if(ISTGIF_CONTAINS("Station"       )) fStrawTrkGeom->fNumberOfStations++;

  // Skip the other StrECAL components
  if     (ISTGIF_CONTAINS("ECAL"          ) ||
          ISTGIF_CONTAINS("StrECALTarget" )) ret = false;
#undef ISTGIF_CONTAINS
    
  COMETNamedVerbose("IStrawTrkGeomVisitor",
                    "IStrawTrkGeomVisitor::VisitNode, " << name
                    << " : " << (ret? "true":"false") );
  return ret;
}

// -------------------------------------------------------------------------------------
int COMET::IStrawTrkGeom::GetSequenceIdClosestToGlobalPosition(const TVector3& global,
                                                               bool findOnSubLayer) const
{
  // In a straw tube region?
  int sequenceId = GetSequenceIdFromGlobalPosition(global);
  if(0 <= sequenceId) return sequenceId;

  // ----------------------------------------------------
  // If not in a straw tube, check close region
  // ----------------------------------------------------
  double strawRadius = 0.5 * GetStrawDiameter();
  double dist, minDist = 1e30;
  TVector3 diffPos, wdir, vdist, bdir;
  // Beam direction
  COMET::IGeomInfo::DetectorSolenoid().GetDetVectorInGlobalCoordinate(TVector3(0, 0, 1), bdir);
  // Find surrounding wires and check the distance to each
  int id, reachStep = 0;
  do{
    reachStep++;
    for(int ix=-reachStep; ix<=reachStep; ix+=2*reachStep){
      for(int iy=-reachStep; iy<=reachStep; iy+=2*reachStep){
        for(int iz=-reachStep; iz<=reachStep; iz+=2*reachStep){
          diffPos.SetXYZ(ix*strawRadius, iy*strawRadius, iz*strawRadius);
          // If required to seach on the same sublayer, suppress beam axis component
          if(findOnSubLayer) diffPos -= bdir * diffPos.Dot(bdir);
          // Get the corresponding sequence Id
          id = GetSequenceIdFromGlobalPosition(global+diffPos);
          if(0 > id) continue;

          wdir  = GetWireDirection(id).Unit();
          vdist = global - GetWirePosition(id);
          vdist -= wdir * vdist.Dot(wdir); // Remove the component on the wire axis
          // Compare distance
          dist  = vdist.Mag(); // distance
          if(minDist > dist){
            minDist = dist;
            sequenceId = id;
          }
        }
      }
    }
  } while(0 > sequenceId && 10 > reachStep);
  return sequenceId;
}

// --------------------------------------------------------------------------------
double COMET::IStrawTrkGeom::EstimateDCA(int sequenceId, const TVector3& pos,
                                         const TVector3& mom, int* signLR, TVector3* dcaPoint) const
{
  if(0 > sequenceId || GetNumberOfChannels() <= sequenceId){
    COMETNamedError("IStrawTrkGeom", "GetMinimumDistanceFromWireOfTrack() : "
                    "Invalid sequenceId or geometryId was input.");
    return 1e30;
  }
  // Get correct wire direction (unit vector)
  TVector3 wireDir = GetWireDirection(sequenceId);
  // Get momentum direction in unit vector
  TVector3 momDir  = mom.Unit();
  // Dot product of both direction
  Double_t dotDirs = wireDir.Dot(momDir);

  // Calc vector representing DCA from a point on the wire (A) to a point on the track (B)
  // Position vectors of A and B are, with scalar values a and b,
  //   P_A = P_wire + a * wireDir
  //   P_B = P_pos  + b * momDir
  // The values a and b are determined so that the DCA vector
  // becomes perpendicular to the wire/momentum direction : 
  //   (P_B - P_A) .dot. wireDir = 0
  //   (P_B - P_A) .dot. momDir  = 0
  
  // Get relative position vector of 'pos' from the wire position : P_pos - P_wire
  TVector3 vpos = pos - GetWirePosition(sequenceId);
  // Solve values a and b
  Double_t a    = vpos.Dot( wireDir - dotDirs * momDir ) / ( 1 - dotDirs*dotDirs );
  Double_t b    = vpos.Dot( dotDirs * wireDir - momDir ) / ( 1 - dotDirs*dotDirs );
  // Get DCA vector = P_B - P_A
  TVector3 dcaVect = vpos + b*momDir - a*wireDir;

  // Give the sign L/R
  if(signLR){
    // According to GFWireHitPolicy, 
    // A vector to define the RIGHT side for these wire/momentum directions
    // i.e. measurment axis for genfit : track dir x wire dir
    // Genfit definition: plane(w) = measurement(u) x wire(v)
    TVector3 n = momDir.Cross(wireDir);
    // Project vpos on the measurment axis and check its sign
    Double_t val = vpos.Dot(n);
    if(0.0 == val) *signLR = 0; // at center
    else           *signLR = (0 < val)? +1 : -1;
  }
  // Give the DCA point = P_B
  if(dcaPoint) *dcaPoint = pos + b*momDir;
  
  return dcaVect.Mag();
}

double COMET::IStrawTrkGeom::EstimateDCA(const COMET::IGeometryId& id, const TVector3& pos,
                                         const TVector3& mom, int* signLR, TVector3* dcaPoint) const
{
  return EstimateDCA(GetSequenceIdFromGeomId(id), pos, mom, signLR, dcaPoint);
}