//
// fana.C - C++ utility classes for the fana pyroot macro
//

#include <fana.h>
#include <cmath>
#include <complex>
#include <TROOT.h>
#include <TTree.h>
#include <TTreeFormula.h>
#include <TH1D.h>
#include <TFile.h>

const double F1_clock_period = 32. / 152 * 128 / 232;
const double F1_rollover_period = 3744;

unsigned int Fana::SetPoints(TChain *ttagm, const char *varexp,
                                            const char *cond,
                                            int maxentries)
{
   /// Load time series vector with values returned by expression varexp
   /// evaluated against the rows of tree ttagm that satisfy conditional
   /// expression cond.

   fVarexp = std::string(varexp);
   TTreeFormula varform("varform", varexp, ttagm);
   const char defcond[] = "1";
   if (cond == 0 || strlen(cond) == 0)
      cond = defcond;
   TTreeFormula condform("condform", cond, ttagm);

   fTseries.clear();
   int maxentry = (maxentries > 0)? maxentries : ttagm->GetEntries();
   double eventtime;
   double lasttime = 0;
   TBranch *branch = 0;
   int tree_size = 0;
   int i0 = 0;
   double var;
   for (int i=0; i < maxentry; ++i) {
      ttagm->LoadTree(i);
      if (i == i0 + tree_size) {
         i0 = i;
         ttagm->SetBranchAddress("eventtime", &eventtime, &branch);
         condform.UpdateFormulaLeaves();
         varform.UpdateFormulaLeaves();
         tree_size = ttagm->GetTree()->GetEntries();
      }
      branch->GetEntry(i-i0);
      if (cond == defcond || condform.EvalInstance() != 0) {
         int repeats;
         if (eventtime != lasttime) {
            fTseries.push_back(varform.EvalInstance());
            lasttime = eventtime;
            repeats = 0;
         }
         else {
            if (repeats > 10000) {
               printf("Error - TChain lost sync, too many repeated values!\n");
               return 0;
            }
            ++repeats;
         }
      }
   }
   ttagm->ResetBranchAddresses();
   return fTseries.size();
}

TH1D *Fana::GetTransform(int n, double *ftmag) {
   /// Compute the Fourier transform of fTseries at regularly-spaced
   /// frequencies passed in the n elements of argument ftmag. Return
   /// the results as a frequency histogram containing the magnitudes
   /// of the Fourier coefficients as the bin contents.

   double f0 = ftmag[0];
   double f1 = ftmag[1] - ftmag[0];
   std::vector<std::complex<double> > phase;
   std::vector<std::complex<double> > efactor;
   const std::complex<double> i2pi(0, 2*M_PI);
   for (unsigned int i=0; i < fTseries.size(); ++i) {
      efactor.push_back(exp(i2pi * fmod(fTseries[i] * f0, 1)));
      phase.push_back(exp(i2pi * fmod(fTseries[i] * f1, 1)));
   }
   std::vector<std::complex<double> > ft;
   for (int p=0; p < n; ++p) {
      std::complex<double> sum;
      for (unsigned int i=0; i < fTseries.size(); ++i) {
         sum += efactor[i];
#if VERBOSE
         printf("p=%d,i=%d,efactor[i]=(%f,%f),sum=(%f,%f)\n",
                p,i,efactor[i].real(),efactor[i].imag(),sum.real(),sum.imag());
#endif
         efactor[i] *= phase[i];
      }
      ft.push_back(sum);
   }

   std::string title("Fourier transform of ");
   title += fVarexp;
   TH1D *fthist = new TH1D("ft", title.c_str(), n, f0, f0 + n * f1);
   for (int p=0; p < n; ++p)
      fthist->SetBinContent(p+1, abs(ft[p]));
   return fthist;
}

TH1I *Fana::Check_t0(TChain *ttagm, const char *cond, int maxentries) {
   /// Scan through a sequence of events in the ttagm tree and histogram
   /// the tdc times of all hits in the TAGM detector relative to the time
   /// of the event trigger. The tdc returns the value of a free-running
   /// counter that gets reset and synchronized to the trigger clock at the
   /// beginning of a run. The tdc clock is supposed to be phase-locked to
   /// the trigger clock, so taking the hit-trigger time difference should
   /// show the fixed boundaries ot the trigger window set in the tdc
   /// firmware. If these boundaries are poorly defined or walk with time,
   /// it shows the drift between the tdc and trigger clocks.

   TH1I *h = dynamic_cast<TH1I*>(gROOT->FindObject("checkt0"));
   if (h != 0)
      h->Delete();
   h = new TH1I("checkt0", "check of thit - ttrig (ns)", 2000, -500., 500.);
   h->SetBit(TH1::kCanRebin);

   ttagm->LoadTree(0);
   const char defcond[] = "1";
   if (cond == 0 || strlen(cond) == 0)
      cond = defcond;
   TTreeFormula condform("condform", cond, ttagm);

   int maxentry = (maxentries > 0)? maxentries : ttagm->GetEntries();
   
   int tdctick;
   double trigtime;
   TBranch *b_tdctick;
   TBranch *b_trigtime;
   int tree_size = 0;
   int i0 = 0;
   for (int i=0; i < maxentry; ++i) {
      ttagm->LoadTree(i);
      if (i == i0 + tree_size) {
         i0 = i;
         condform.UpdateFormulaLeaves();
         ttagm->SetBranchAddress("tick", &tdctick, &b_tdctick);
         ttagm->SetBranchAddress("trigtime", &trigtime, &b_trigtime);
         tree_size = ttagm->GetTree()->GetEntries();
      }
      b_tdctick->GetEntry(i-i0);
      b_trigtime->GetEntry(i-i0);
      if (cond != defcond && condform.EvalInstance() == 0)
         continue;
      double ttrigger = fmod(trigtime, F1_rollover_period);
      double tphase = tdctick * F1_clock_period/2 - ttrigger;
      tphase += (tphase > -1000)? 0 : F1_rollover_period;
      h->Fill(tphase);
   }
   ttagm->ResetBranchAddresses();
   return h;
}

TH1I *Fana::Check_rf(TChain *ttagm, double rf_freq, int maxentries) {
   /// Scan through a sequence of events in the ttagm tree and histogram
   /// the prescaled rf tdc time stored in each event relative to what it
   /// would have been if the prescaled rf frequency were rf_freq and the
   /// clocks were perfect. The so-called rf phase (ns) is a running sum
   /// of these offsets. It gets computed and stored in an output TTree
   /// that is "friend" to the original ttagm tree. The returned histogram
   /// contains the distribution of the phases for all events scanned.

   TH1I *h = dynamic_cast<TH1I*>(gROOT->FindObject("checkrf"));
   if (h != 0)
      h->Delete();
   h = new TH1I("checkrf", "rf - tdc clock phase distribution", 200, -1., 1.);
   h->SetBit(TH1::kCanRebin);

   TTree *tphase = ttagm->GetFriend("tphase");
   if (tphase != 0)
      ttagm->RemoveFriend(tphase);
   TFile froot("fana.root", "UPDATE");
   tphase = new TTree("tphase", "rf phase distribution (ns)");
   double rfphase;
   tphase->Branch("rfphase", &rfphase, "rfphase/D");

   int maxentry = (maxentries > 0)? maxentries : ttagm->GetEntries();
   
   double rf_freq0 = rf_freq;
   rfphase = 0;
   double rftime = 0;
   double rftlast = 0;
   TBranch *branch;
   int tree_size = 0;
   int i0 = 0;
   for (int i=0; i < maxentry; ++i) {
      ttagm->LoadTree(i);
      if (i == i0 + tree_size) {
         i0 = i;
         ttagm->SetBranchAddress("rftime", &rftime, &branch);
         tree_size = ttagm->GetTree()->GetEntries();
      }
      branch->GetEntry(i-i0);
      double taurf = 1 / rf_freq;
      double dt = rftime - rftlast;
      double delta = fmod(dt + 999999.5 * taurf, taurf) - 0.5 * taurf;
      if (rftime > 0 && (fabs(delta) < fabs(dt) * 1e-6 || rftlast == 0)) {
         double fratio = rf_freq / rf_freq0;
         double dphase = (fratio - 1) * dt - delta * fratio;
         rfphase += fmod(dphase + 999999.5 * taurf, taurf) - 0.5 * taurf;
         rf_freq += -rf_freq * (delta / dt) * 0.01;
         rftlast = rftime;
      }
      h->Fill(rfphase);
      tphase->Fill();
   }
   tphase->Write();
   froot.Close();
   ttagm->ResetBranchAddresses();
   ttagm->AddFriend("tphase", "fana.root");
   return h;
}