#include <sstream>
#include <string>
#include <cmath>
#include <memory>

#include <TVector3.h>
#include <TMatrixD.h>
#include <TPrincipal.h>

#include <HEPUnits.hxx>
#include <IHitSelection.hxx>
#include <IRealDatum.hxx>
#include <IReconNode.hxx>
#include <IReconState.hxx>
#include <IReconTrack.hxx>
#include <IHitFilter.hxx>
#include <ICOMETLog.hxx>
#include "ILineFit.hxx"
#include "createIReconTrack.hxx"
#include "IReconTrack.hxx"
#include "IReconCluster.hxx"
#include "IReconNode.hxx"


// This method creates a ITrackState from a position 
// and direction (with covariance matrices) and a IHit.
COMET::IHandle<COMET::ITrackState> createTrackState(const TVector3& pos,
                const TMatrixD& posCov,
                const TVector3& dir,
                const TMatrixD& dirCov,
                COMET::IHandle<COMET::IHit> hit,
                bool xValid,
                bool yValid,
                bool zValid) {

  COMET::IHandle<COMET::ITrackState> tstate(new COMET::ITrackState);

  // Offset in IReconState
  int pi = 1;
  int di = 5;
  int ki = 8;

  // Set the EDeposit
  tstate->SetValue(0,hit->GetCharge());
  tstate->SetCovarianceValue(0,0,0);

  // Set the value and covariance matrix for the position
  for(int i = 0;i < 3;i++){
    tstate->SetValue(pi+i,pos[i]);      
    for(int j = 0; j < 3; j++)
       tstate->SetCovarianceValue(pi+i,pi+j,posCov(i,j));      
  }
  
  // Set the time
  tstate->SetValue(4,hit->GetTime());
  tstate->SetCovarianceValue(4,0,0);

  // Get the value and covariance matrix for the direction
  for(int i = 0;i < 3;i++){
    tstate->SetValue(di+i,dir[i]);      
    for(int j = 0; j < 3; j++)
       tstate->SetCovarianceValue(di+i,di+j,dirCov(i,j));      
  }

  // Set the curvature.  It should be zero and fixed, since we assume
  // straight line.
  tstate->SetValue(ki,0.0);
  tstate->SetFixed(ki);

  // Deal with free dimensions
  if(!xValid){
    tstate->SetFree(1);
    tstate->SetFree(5);
  }
  if(!yValid){
    tstate->SetFree(2);
    tstate->SetFree(6);
  }
  if(!zValid){
    tstate->SetFree(3);
    tstate->SetFree(7);
  }
    
  return tstate;

}


namespace {
    TPrincipal* MakePCA(const COMET::IHitSelection& hits,
                        bool& xValid, bool& yValid, bool& zValid) {
        if (hits.size()<3) {
            COMETWarn("No Hits");
            xValid = false;
            yValid = false;
            zValid = false;
            return NULL;
        }
        
        xValid = true;
        yValid = true;
        zValid = true;
        
        TPrincipal* pca = new TPrincipal(3,"");
        
        // Find the charge of the minimum charge hit, but protect against
        // really low charge hits.
        double minQ = 1;
        for (COMET::IHitSelection::const_iterator h = hits.begin(); 
             h != hits.end();
             ++h) {
            minQ = std::min(minQ,(*h)->GetCharge());
            if (!(*h)->IsXHit()) xValid = false;
            if (!(*h)->IsYHit()) yValid = false;
            if (!(*h)->IsZHit()) zValid = false;
        }
        minQ = std::max(0.1,minQ);
        
        for (COMET::IHitSelection::const_iterator h = hits.begin(); 
             h != hits.end();
             ++h) {
            double row[3] = {(*h)->GetPosition().X(),
                             (*h)->GetPosition().Y(),
                             (*h)->GetPosition().Z()};
            for (double q=0; q<(*h)->GetCharge(); q += minQ) {
                pca->AddRow(row);
            }
        }
        pca->MakePrincipals();
        
        return pca;
    }

    class IPCAComparator {
    public:
        IPCAComparator(TPrincipal& pca) : fPCA(pca) {}

        bool operator() (const COMET::IHandle<COMET::IHit>& a, const COMET::IHandle<COMET::IHit>& b) {
            double xA[3] = {a->GetPosition().X(), 
                            a->GetPosition().Y(), 
                            a->GetPosition().Z()};
            double pA[3];
            fPCA.X2P(xA,pA);
            double xB[3] = {b->GetPosition().X(), 
                            b->GetPosition().Y(), 
                            b->GetPosition().Z()};
            double pB[3];
            fPCA.X2P(xB,pB);
            return (pA[0] < pB[0]);
        }
    private:
        TPrincipal& fPCA;
    };

}

COMET::IHandle<COMET::IReconTrack> CreateTrack::createPCAReconTrack(const COMET::IHitSelection& hits) {
    
    if (hits.size()<3) {
        COMETWarn("No Hits");
        return COMET::IHandle<COMET::IReconTrack>();
    }

    bool xValid, yValid, zValid;
    std::auto_ptr<TPrincipal> pca(MakePCA(hits,xValid,yValid,zValid));

    // Yes I know boolean algebra... This expresion could be simplified, but I
    // want to make the intention clear.  This must be an xz or an yz track.
    if ( !((xValid&&zValid) 
           || (yValid&&zValid)
           || (xValid&&yValid)) ) {
        COMETWarn("No Valid PCA Tracks");
        return COMET::IHandle<COMET::IReconTrack>();
    }

    // We have a real track, so create the handle.
    COMET::IHandle<COMET::IReconTrack> track(new COMET::IReconTrack());

    track->SetQuality(0.0); 
    int ndof = hits.size();
    if (zValid) ndof--;
    if (yValid) ndof--;
    if (xValid) ndof--;    
    track->SetNDOF(ndof); 
    track->SetStatus(COMET::IReconTrack::kSuccess); 
    track->SetAlgorithmName("PCATrack");

    // Find the total charge in the shower cluster.
    double totalCharge = 0.0;
    for (COMET::IHitSelection::const_iterator h = hits.begin(); 
         h != hits.end();
         ++h) {
        totalCharge += (*h)->GetCharge();
    }

    // Save the total charge in the track.  
    COMET::IRealDatum* chargeDatum = new COMET::IRealDatum("totalCharge",
                                            "Shower Total Charge");
    chargeDatum->SetValue(totalCharge);
    track->AddDatum(chargeDatum);

    // Save the PCA eigenvalues in the track.
    TVectorD eigenValues(*(pca->GetEigenValues()));
    COMET::IRealDatum* eigenDatum = new COMET::IRealDatum("eigenValues",
                                            "Shower Eigen Values");
    eigenDatum->SetValue(eigenValues(0));
    eigenDatum->GetVector().push_back(eigenValues(1));
    eigenDatum->GetVector().push_back(eigenValues(2));
    track->AddDatum(eigenDatum);


    // Loop through the principle component value for each hit and find the
    // maximum and minimum values.  These will be used to determine the length
    // of the shower as well as the front and back position of the shower.
    double minP = 0;
    double maxP = 0;

    for (COMET::IHitSelection::const_iterator h = hits.begin(); 
         h != hits.end();
         ++h) {
        double x[3] = {(*h)->GetPosition().X(), 
                       (*h)->GetPosition().Y(), 
                       (*h)->GetPosition().Z()};
        double p[3];
        pca->X2P(x,p);
        minP = std::min(minP,p[0]);
        maxP = std::max(maxP,p[0]);

    }

    // Translate the minimum principle component into the front position.
    TVector3 frontPos;
    double p[3], x[3];
    p[0] = minP; p[1] = 0; p[2] = 0;
    pca->P2X(p,x,3);
    frontPos.SetX(x[0]); frontPos.SetY(x[1]); frontPos.SetZ(x[2]);

    // Translate the maximum principle component into the back position.
    TVector3 backPos;
    p[0] = maxP; p[1] = 0; p[2] = 0;
    pca->P2X(p,x,3);
    backPos.SetX(x[0]); backPos.SetY(x[1]); backPos.SetZ(x[2]);

    // Flip the front and back points if necessary in order to make
    // front point the upstream point.
    bool flip = false;
    if(frontPos.Z()> backPos.Z()){
      flip = true;
      TVector3 tmp = frontPos;
      frontPos = backPos;
      backPos = tmp;
    }

    // Get the position covariance from the pca covariance matrix.  
    TMatrixD posCov(*(pca->GetCovarianceMatrix()));
    posCov = posCov * (1.0/totalCharge);

    // Estimate the direction of the shower based on the calculated front and
    // back positions.  This needs to have a direction covariance, but for
    // now... it doesn't exist.  Use a zero matrix.
    TVector3 diff = 0.5*(backPos - frontPos);
    TVector3 dir = diff.Unit();
    TMatrixD dirCov(3,3);


    // Now put list of hits in correct order.
    COMET::IHitSelection *new_hits = new COMET::IHitSelection(hits);

    std::sort(new_hits->begin(), new_hits->end(),
            IPCAComparator(*pca));

    // Need to flip the order for some tracks... 
    // Do this before calculating PCA for each hit...
    if(flip)
      std::reverse(new_hits->begin(), new_hits->end());

    //COMET::IReconNodeContainer nodes;
    COMET::IReconNodeContainerImpl<COMET::ITrackState> nodes;

    for (COMET::IHitSelection::const_iterator h = new_hits->begin(); 
         h != new_hits->end(); ++h) {


        // Okay, need to play a little game here.  
        // Calculate the principal components for the hit position.
        // Then use only the most significant component value 
        // in order to calculate the X,Y,Z of the PCA line 
        // for this hit position.
        double x[3] = {(*h)->GetPosition().X(), 
                       (*h)->GetPosition().Y(), 
                       (*h)->GetPosition().Z()};
        double p[3];
        pca->X2P(x,p);      // Convert hit position to principal components.
        p[1] = 0; p[2] = 0; // Set the less significant components to zero.
        pca->P2X(p,x,3);    // Calculated the 'fitted' position for this hit.
        TVector3 pos(x[0],x[1],x[2]); 
        

        COMET::IHandle<COMET::IReconNode> node(new COMET::IReconNode);
  //        COMET::IHandle<COMET::IReconBase> nodeObject(new COMET::IReconBase);
  COMET::IHandle<COMET::IReconBase> nodeObject(new COMET::IReconCluster);
        COMET::IHitSelection* nodeHits = new COMET::IHitSelection("nodeHits");
        nodeHits->push_back(*h);
        nodeObject->AddHits(nodeHits);
        node->SetObject(nodeObject);

  COMET::IHandle<COMET::ITrackState> nodeState = createTrackState(pos,posCov,
                  dir,dirCov,(*h),
                  xValid,yValid,zValid);
        
        COMET::IHandle<COMET::IReconState> castState(nodeState);
        node->SetState(castState);
        nodes.push_back(node);

    }

    // Finally, define the overall state of the track using the first node state.
    COMET::IHandle<COMET::ITrackState> fittedState = track->GetState();
    COMET::IHandle<COMET::ITrackState> nodeState = nodes.front()->GetState();
        
    if (fittedState && nodeState) {
        *fittedState = *nodeState;
    }

    // Add hits to track.
    track->AddHits(new_hits);

    // Add nodes to track.
    std::copy(nodes.begin(), nodes.end(),
              std::back_inserter(track->GetNodes()));

    return track;
}



COMET::IHandle<COMET::IReconTrack> CreateTrack::createReconTrack(const COMET::IHitSelection& hits) {
    
    if (hits.size()<3) {
        COMETWarn("No Hits");
        return COMET::IHandle<COMET::IReconTrack>();
    }

    // We have a real track, so create the handle.
    COMET::IHandle<COMET::IReconTrack> track(new COMET::IReconTrack());

    track->SetQuality(0.0); 
    int ndof = hits.size();
    track->SetNDOF(ndof); 
    track->SetStatus(COMET::IReconTrack::kSuccess); 
    track->SetAlgorithmName("PCATrack");

    // Find the total charge in the shower cluster.
    double totalCharge = 0.0;
    for (COMET::IHitSelection::const_iterator h = hits.begin(); 
         h != hits.end();
         ++h) {
        totalCharge += (*h)->GetCharge();
    }

    // Save the total charge in the track.  
    COMET::IRealDatum* chargeDatum = new COMET::IRealDatum("totalCharge",
                                            "Shower Total Charge");
    chargeDatum->SetValue(totalCharge);
    track->AddDatum(chargeDatum);

    // Now put list of hits in correct order.
    COMET::IHitSelection *new_hits = new COMET::IHitSelection(hits);

    // Add hits to track.
    track->AddHits(new_hits);

    // Add nodes to track.
    //    std::copy(nodes.begin(), nodes.end(),
    //          std::back_inserter(track->GetNodes()));

    return track;
}