#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 ang_rad, float momentum, int ispion=0)
{
  int inside = 0;
  float ang_deg = ang_rad*180./TMath::Pi();
  float radius = 0.26;
  float length = 0.3;
  float pass_l = 0.03; // decay length of lambda: 3 cm
  float mom_th = radius*2.*3.e8*1.e-9;
  float p_l = momentum*cos(ang_rad);
  float p_p = momentum*sin(ang_rad);
  if(p_p>=mom_th)
    if(p_l == 0) inside = 1; 
    else if(ang_deg>=90)
      {
	float tt = -(radius+(float)ispion*pass_l)/p_l;
	float rr = p_p*1.e9/2./3.e8;
	float ang = p_p*tt/rr;
	float dist = 2*rr*sin(ang/2.);
	if(dist >= radius || ang >= TMath::Pi()/2.) inside = 1;
      }
    else
      {
	float tt = (length-(float)ispion*pass_l)/p_l;
	float rr = p_p*1.e9/2./3.e8;
	float ang = p_p*tt/rr;
	float dist = 2*rr*sin(ang/2.);
	if(dist >= radius || ang >= TMath::Pi()/2.) inside = 1;
      }
  return inside;
}

void filter_pkk(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,"b")) 
    {
      flag_p2 = 1;
      printf("Filtering proton 2 as well.\n");
    }
  
  // define variables for new Tree
  int det_k1, det_k2, det_p1, det_p2;  // k1 is K+, k2 is K-
  
  // define variables for old Tree
  float p_k1, th_k1, ph_k1;
  float p_k2, th_k2, ph_k2;
  float p_p1, th_p1, ph_p1;
  float p_p2, th_p2, ph_p2;

  // 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("k1.af.P",&p_k1);
  T1->SetBranchAddress("k1.af.Th",&th_k1);
  T1->SetBranchAddress("k1.af.Ph",&ph_k1);

  T1->SetBranchAddress("k2.af.P",&p_k2);
  T1->SetBranchAddress("k2.af.Th",&th_k2);
  T1->SetBranchAddress("k2.af.Ph",&ph_k2);

  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("k1.det",&det_k1,"k1.det/I");
  T2->Branch("k2.det",&det_k2,"k2.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_k1 = det_k2 = 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.5;
      float p_min_pi = 0.15;
      float p_min_p = 0.7;

      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_k1)&&th_k1>ang_min_p&&th_k1<ang_max_p&&p_k1>p_min_k)
	det_k1++; 
      
      // kaon - 
      if(InSector(ph_k2)&&th_k2>ang_min_n&&th_k2<ang_max_n&&p_k2>p_min_k)
	det_k2++; 
      
      // 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, p_p1)&&p_p1<tof_p_lim) det_p1++;
      // kaon +
      if(TouchRing(th_k1, p_k1)&&p_k1<tof_k_lim) det_k1++;
      // kaon -
      if(TouchRing(th_k2, p_k2)&&p_k2<tof_k_lim) det_k2++;
      // proton2
      if(flag_p2)
	if(TouchRing(th_p2, p_p2)&&p_p2<tof_p_lim) det_p2++;


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