#include<cstdio>
#include<TTree.h>
#include<TSystem.h>
#include<TLorentzVector.h>
#include<TStopwatch.h>
#include<TRandom.h>
#include<TGenPhaseSpace.h>
#include<TApplication.h>
#include<TFile.h>
#include<TMath.h>
#include<TROOT.h>

int InSector(float ang_rad)
{
  float ang_deg = ang_rad*180./TMath::Pi();
  int inside = 1;
  for(int i = -6; i<6; i++)
    {
      if(fabs(ang_deg-i*60)<7.5) inside = 0;
    }
  return inside;
}

int TouchRing(float theta_rad, float phi_rad, float momentum, int charge=-1)
{
  int in_accp = 0;
  float d_i = 0.5;   // distance before magnet
  float d_m = 1.;    // magnet length
  float d_f = 1.;    // distrance after magnet
  float open_h = 0.55;   // vertical opening
  float open_w = 1.35;   // horizontal opening

  float ang_in = atan(tan(theta_rad)*cos(phi_rad));
  float ang_out = atan(tan(theta_rad)*sin(phi_rad));
  float ang_out_cut = open_h/(d_i+d_m)/2.;
  
  if( (fabs(ang_out)>ang_out_cut) || (theta_rad>1) ) in_accp = 0;
  else
    {
      float ang_tmp = ang_in*(float)charge;
      float bb = 0.7; // field strength
      float rr = momentum/0.3/bb; // radius
      if( rr*(1+sin(ang_tmp)) <= d_m ) in_accp = 0;
      else
	{
	  float x_i = d_i*tan(ang_tmp); // entrance position
	  float ang_bend = - asin((d_m - rr*sin(ang_tmp))/rr);
	  float x_m = x_i + rr*(cos(ang_bend)-cos(ang_tmp));
	  if( fabs(x_m) > open_w/2.) in_accp = 0;
	  else
	    {
	      float x_f = x_m + d_f*tan(ang_bend);
	      if( fabs(x_f) > 2.5 ) in_accp = 0;
	      else in_accp = 1;
	    }
	}
    }
  return in_accp;
}
  
void filter(char* t_name)
{
  
  // open root tree
  TFile* ff1 = new TFile("tree_cu63.root","r");
  ff1->cd();

  char tree_name[20], tree_name_addon[20];
  sprintf(tree_name, "T_%s", t_name);
  sprintf(tree_name_addon, "T_%s_addon", t_name);
  TTree* T1 = (TTree*)gROOT->FindObject(tree_name);
  if(!T1) return;

  TFile* ff2 = new TFile("tree_cu63_addon.root","update");
  ff2->cd();
  
  // flag for p2 existance
  int flag_p2 = 0;
  printf("tree name is %s\n", t_name);
  if(!strcmp(t_name,"bl")||!strcmp(t_name,"bs")) 
    {
      flag_p2 = 1;
      printf("Filtering proton 2 as well.\n");
    }
  
  // define variables for new Tree
  int det_pi, det_k, det_p1, det_p2;
  
  // define variables for old Tree
  float p_k, th_k, ph_k;
  float p_p1, th_p1, ph_p1;
  float p_p2, th_p2, ph_p2;
  float p_pi, th_pi, ph_pi;
  //  float p_g, th_g, ph_g;

  // Project variables from old Tree

  T1->SetBranchAddress("p1.af.P",&p_p1);
  T1->SetBranchAddress("p1.af.Th",&th_p1);
  T1->SetBranchAddress("p1.af.Ph",&ph_p1);

  T1->SetBranchAddress("k.af.P",&p_k);
  T1->SetBranchAddress("k.af.Th",&th_k);
  T1->SetBranchAddress("k.af.Ph",&ph_k);

  T1->SetBranchAddress("pi.af.P",&p_pi);
  T1->SetBranchAddress("pi.af.Th",&th_pi);
  T1->SetBranchAddress("pi.af.Ph",&ph_pi);

  if(flag_p2)
    {
      T1->SetBranchAddress("p2.af.P",&p_p2);
      T1->SetBranchAddress("p2.af.Th",&th_p2);
      T1->SetBranchAddress("p2.af.Ph",&ph_p2);
    }

  // Create new Tree to store acceptance flag
  TTree* T2 = new TTree(tree_name_addon,"Tree for Additional Information");
  T2->Branch("p1.det",&det_p1,"p1.det/I");
  T2->Branch("pi.det",&det_pi,"pi.det/I");
  T2->Branch("k.det",&det_k,"k.det/I");
  if(flag_p2) T2->Branch("p2.det",&det_p2,"p2.det/I");

  // Major loop
  int nn = T1->GetEntries();
  for (int i = 0; i<nn; i++)
    { 
      det_pi = det_k = det_p1 = det_p2 = 0;
      if((i+1)%50000==0) printf("%d events processed...\n",i+1);
      T1->GetEntry(i);

//       // Forward detector
//       float ang_min_p = 5.*TMath::Pi()/180.;
//       float ang_max_p = 35.*TMath::Pi()/180.;
//       float ang_min_n = 8.*TMath::Pi()/180.;
//       float ang_max_n = 40.*TMath::Pi()/180.;

      float p_min_k = 0.0;
      float p_min_pi = 0.0;
      float p_min_p = 0.0;

//       float tof_k_lim = 0.65;
//       float tof_pi_lim = 0.65;
//       float tof_p_lim = 1.03;

//       // proton 1
//       if(InSector(ph_p1)&&th_p1>ang_min_p&&th_p1<ang_max_p&&p_p1>p_min_p)
// 	det_p1++; 
      
//       // kaon +
//       if(InSector(ph_k)&&th_k>ang_min_p&&th_k<ang_max_p&&p_k>p_min_k)
// 	det_k++; 
      
//       // pion - 
//       if(InSector(ph_pi)&&th_pi>ang_min_n&&th_pi<ang_max_n&&p_pi>p_min_pi)
// 	det_pi++; 
      
//       // proton 2
//       if(flag_p2)
// 	if(InSector(ph_p2)&&th_p2>ang_min_p&&th_p2<ang_max_p&&p_p2>p_min_p)
// 	  det_p2++; 
      
      // Central detector
      //      proton1
      if(TouchRing(th_p1, ph_p1, p_p1, 1)&&p_p1>p_min_p) det_p1++;
      // kaon +
      if(TouchRing(th_k, ph_k, p_k, 1 )&&p_k>p_min_k) det_k++;
      //      pion -
      if(TouchRing(th_pi, ph_pi, p_pi, -1)&&p_pi>p_min_pi) det_pi++;
      //      proton2
      if(flag_p2)
	if(TouchRing(th_p2, ph_p2, p_p2, 1)&&p_p2>p_min_p) det_p2++;


      T2->Fill();
      
    }
  ff2->cd();
  T2->Write("",TObject::kOverwrite);
  ff2->Close();
  ff1->Close();
}