#include "TF1.h"
#include "Stntuple/loop/TStnAna.hh"
#include "Stntuple/fit/TCtvmft.hh"
#include "TPad.h"
#include "TCanvas.h"

#include "TKShortCtvmftModule.hh"

#include "TLorentzVector.h"
#include <iostream>
#include <vector>

bool trackQuality(TStnTrack* t) {
  bool ret = false;
  if (t->NCotHitsAx()>=20
      && t->NCotHitsSt()>=20)
    ret = true;
  return ret;
}

//_____________________________________________________________________________
TKShortCtvmftModule::TKShortCtvmftModule(const char* name, const char* title):
  TStnModule(name,title)
{
  fKShortFit = new TCtvmft();
}

//_____________________________________________________________________________
TKShortCtvmftModule::~TKShortCtvmftModule() {
  delete fKShortFit;
}


//_____________________________________________________________________________
void TKShortCtvmftModule::BookHistograms() {
  char name [200];
  char title[200];

  Delete("hist");
                                // book histograms

  cout << "Booking Histograms" << endl;

  sprintf(name, "%s_KShortMass",GetName());
  sprintf(title,"%s: KShort Mass",GetName());
  fKShortMass = new TH1F(name, title,100,0.2,0.7);
  AddHistogram(fKShortMass);

  sprintf(name, "%s_KShortLxy",GetName());
  sprintf(title,"%s: KShort Lxy",GetName());
  fKShortLxy = new TH1F(name, title,100,-1.,1.);
  AddHistogram(fKShortLxy);
}

//_____________________________________________________________________________
int TKShortCtvmftModule::BeginJob() {

                                // register the data block
  cout << "BeginJob" << endl;
  fTrackBlock = (TStnTrackBlock*) fAna->RegisterDataBlock("TrackBlock","TStnTrackBlock");

                                        // book histograms
  BookHistograms();

  if (! fTrackBlock) {
    printf(" >>> branch *** %s *** doesn't exist \n","TrackBlock");
    fEnabled = 0;
  }

  return 0;
}


//_____________________________________________________________________________
int TKShortCtvmftModule::BeginRun() {
  cout << "BeginRun" << endl;
  fDbm = TStnDBManager::Instance();

  fbpcot = (TStnBeamPos*) fDbm->GetTable("CotBeamPos");
  fbpsvx = (TStnBeamPos*) fDbm->GetTable("SvxBeamPos");

  printf("CotBeam x=%8.5f y=%8.5f  sx=%10.7f  sy=%10.7f\n",
         fbpcot->X0(),fbpcot->Y0(),fbpcot->DxDz(),fbpcot->DyDz());
  printf("SvxBeam x=%8.5f y=%8.5f  sx=%10.7f  sy=%10.7f\n",
         fbpsvx->X0(),fbpsvx->Y0(),fbpsvx->DxDz(),fbpsvx->DyDz());

  return 0;
}

//_____________________________________________________________________________
int TKShortCtvmftModule::Event(int ientry) {

  fTrackBlock->GetEntry(ientry);
                                        // do whatever you want
  if (PrintLevel()) {
        printf(" ---- ientry = %7i\n",ientry);
    //    fHeaderBlock->Print();
    //    fCmuDataBlock->Print();
  }

  int ntr = fTrackBlock->NTracks();

  const float PI_MASS =  0.13957;

  // primary vertex
  TVector3 pv(fbpsvx->X0(),fbpsvx->Y0());

  for (int i=0; i<ntr-1; i++) {
    TStnTrack* ti = fTrackBlock->Track(i);
    if (trackQuality(ti)) {
      for (int j=i+1; j<ntr; ++j) {
        TStnTrack* tj = fTrackBlock->Track(j);

        if (trackQuality(tj) 
            && ti->Charge()*tj->Charge()==-1.) {

          // Add Pi+, Pi-
          fKShortFit->Init();
          fKShortFit->AddTrack(ti,PI_MASS,1);
          fKShortFit->AddTrack(tj,PI_MASS,1);

          fKShortFit->SetPrimaryVertex(&pv);
          // don't care about the error on the primary vertex
          // in this example
        
          // Do the fit.
          bool status = fKShortFit->Fit();
//        if (!status) {
//          fKShortFit->PrintErr();
//        }
          if (status) {
//          fKShortFit->Print();

            TVector3 vertex;
            fKShortFit->GetVertex(1,vertex);

            TLorentzVector p4PiP, p4PiM;
            fKShortFit->GetTrackP4(ti->Number(), p4PiP);
            fKShortFit->GetTrackP4(tj->Number(), p4PiM);
          
            TLorentzVector p4KShort = p4PiP + p4PiM;

            float dlxy=0;
            TVector3 kShortDir(p4KShort.X(), p4KShort.Y(), 0.);
            kShortDir = kShortDir.Unit();
            float lxy = fKShortFit->GetDecayLength(0, 1, kShortDir, dlxy);
            fKShortLxy->Fill(lxy);
            fKShortMass->Fill(p4KShort.M());

          }
        }
      }
    }
  }
  return 0;                    
}


//_____________________________________________________________________________
int TKShortCtvmftModule::EndJob() {
  printf("----- end job: ---- %s\n",GetName());

  return 0;
}


//_____________________________________________________________________________
void TKShortCtvmftModule::PlotHistograms(int run_number, int slide) {
  // plot slides

  char name[120], canvas_name[120], title[120], pad_name[120];

  sprintf(name,"run_%i_slide_%i",run_number,slide);

  sprintf(title,"run %i slide %i ",run_number, slide);

  sprintf(canvas_name,"%s_%s",GetName(),name);
  sprintf(pad_name,"%s_p1",canvas_name);

  TCanvas* c = NewSlide(name,title,2,3);
  TPad* p1 = (TPad*) c->GetPrimitive(pad_name);
  
  p1->cd(1); 
  fKShortMass->Draw();

  p1->cd(2);
  fKShortLxy->Draw();
  
  gPad->Update();
}