#include<iostream>
#include<cmath>

#include <iostream>
#include <algorithm>
#include "TFile.h"
#include "TTree.h"
#include "TCanvas.h"
#include "TH1F.h"
#include "TClonesArray.h"

const double B_f = 0.22;
const double alpha = 0.0072973525643;
const double alpha_f = 1.0; // for lepton case
const double MP= TMath::Pi();
const double mtop=173.0;
const double s=pow(360,2);
const double mt =173.0; // top quark
const double Mw =80.419; //W mass
const double r =pow(Mw,2)/pow(mt,2);
const double W =pow((1.0-r),2.) * (1.0+(2.*r));

//Define the SM factors
//const double SinthetaW = 0.22;
//sin^2 θW
const double CosthetaW = 0.881903320;
const double SinthetaW =0.2222465;
//Top velocity :
const double beta_val = sqrt((1.0 - (4.0 * pow(mtop,2))/s));
const double B = (1.0-beta_val)/(1.0+beta_val);

const double Ay = (4.0 / 3.0) * sqrt(SinthetaW);
const double By = 0.0;
const double Az = (1.0 / (2.0 * sqrt(CosthetaW))) * (1.0 - (8.0 / 3.0) * SinthetaW);
const double Bz = (1.0 / (2.0 * sqrt(CosthetaW)));

const double ve = -1.0 + (4.0 * SinthetaW);

// Non SM CONTRIB:
const double Dgamma=0.005;
const double Mz = 91.188;
const double d = (s / (s - pow(Mz, 2))) * (1 / (16 * SinthetaW * pow(CosthetaW,2)));
const double d_p = s / (4 * sqrt(SinthetaW) * sqrt(CosthetaW) * (s - pow(Mz,2)));

const double C = 1 / (4 * SinthetaW);

double DV(double Ay, double Az, double d_p, double ve,double C) {
    double term1 = Ay*Ay - 2*Ay*Az*ve*d_p + Az*Az*(1+pow(ve,2))*d_p*d_p;
    double term2 = 2 * (Ay - Az*ve*d_p) * Dgamma;
    double term3 = -2 * (Ay*ve*d_p - Az*(1+pow(ve,2))*pow(d_p,2)) * Dgamma;
    return (C*(term1+term2+term3));
}

double DA(double Bz,double C,double ve, double d_p) {
    double term1 = Bz*Bz*(1+pow(ve,2))*pow(d_p,2);
    double term2 = -2*Bz*ve*d_p*Dgamma;
    double term3 = 2*Bz*(1+pow(ve,2))*pow(d_p,2)*Dgamma;
    return (C*(term1+term2+term3));
}

double DVA(double Ay, double Az,double C,double Bz, double d_p, double ve) {
    double term1 = -Ay*Bz*d_p*ve  + Az*Bz*(1+pow(ve,2))*d_p*d_p - Bz*ve*d_p*Dgamma;
    double term2 = (Ay - ve*d_p*Az) * Dgamma + Bz*(1+ pow(ve,2))*pow(d_p,2)*Dgamma;
    double term3 = -(Ay*ve*d_p - Az*(1+pow(ve,2))*pow(d_p,2)) * Dgamma;
    return (C*(term1+term2+term3));
}

double EV(double Ay, double Az, double C,double d_p,double ve) {
    double term1 = 2*C * (Ay*Az*d_p - pow(Az,2)*ve*d_p*d_p + Az*d_p*Dgamma + (Ay*d_p - 2*Az*ve*pow(d_p,2)) * Dgamma);
    return term1;
}

double EA(double Bz, double ve, double C,double d_p) {
    double term1 = 2*C * (-Bz*Bz*pow(d_p,2)*ve + Bz*d_p*Dgamma - 2*Bz*ve*pow(d_p,2)*Dgamma);
    return term1;
}

double EVA(double Ay, double ve,double Az, double Bz, double C,double d_p) {
    double term1 = Ay*Bz*d_p - 2*Az*Bz*pow(d_p,2)*ve + Bz*d_p*Dgamma + Az*d_p*Dgamma;
    double term2 = -2*Bz*ve*pow(d_p,2)*Dgamma + (Ay*d_p - 2*Az*ve*pow(d_p,2)) * Dgamma;
    return (C*(term1+term2));
}

double H1(double beta_val, double omega_plus, double omega_minus, double x_f, double W) {
  return ( ((1.0-pow(beta_val,2)) / (2.0*W*beta_val*x_f)) *
    ( pow(omega_plus,2)*(3-2*(omega_plus)) - pow(omega_minus, 2)*(3-2*(omega_minus)) ));
}

double H2(double beta_val, double omega_plus, double omega_minus, double x_f, double W) {
    return ( ((1+beta_val) * pow((1.0-beta_val),2) / (4.0*beta_val*W*pow(x_f,2))) *
      ( pow(omega_plus,2)*(6-8*omega_plus+3*pow(omega_plus,2)) -
      pow(omega_minus,2)*(6-8*omega_minus+3*pow(omega_minus,2)) ));
}

double F(double W,double beta_val,double omega_plus, double omega_minus){
  return ( ((1.0+beta_val)/beta_val) * (3.0/(2.0*W)) * (pow(omega_plus,2) - pow(omega_minus,2)) );
}

double G(double W, double beta_val, double omega_plus, double omega_minus, double x_f) {
  return ( F(W,beta_val,omega_plus,omega_minus) + (3.0*pow((1+beta_val),2)/(beta_val*W) *
    x_f * (omega_plus + log(abs(1-omega_plus)) - (omega_minus + log(abs(1-omega_minus))))));
}


double Function_Theta(double x_f, double thetaCosine) {

  double omega_minus=0, omega_plus=0;
  double omega=0;

  if ( r >= B ) {
    if      ( x_f >= B*r && x_f <  B) { omega_minus = 1 - x_f/B; omega_plus  = 1 - r;   }
    else if ( x_f >= B   && x_f <  r) { omega_minus = 0;         omega_plus  = 1 - r;   }
    else if ( x_f >= r   && x_f <= 1) { omega_minus = 0;         omega_plus  = 1 - x_f; }
  }

  if ( r < B ) {
    if      ( x_f >= B*r && x_f <  r) { omega_minus = 1 - x_f/B; omega_plus =  1 - r;   }
    else if ( x_f >= r   && x_f <  B) { omega_minus = 1 - x_f/B; omega_plus =  1 - x_f; }
    else if ( x_f >= B   && x_f <= 1) { omega_minus = 0;         omega_plus =  1 - x_f; }
  }

    double C0_term1 = 0.5 * ( (3-pow(beta_val,2)) * DV(Ay,Az,d_p,ve,C) - (1-3*pow(beta_val,2)) * DA(Bz,C,ve,d_p) );
    double C0_term2 = -alpha_f * (1-pow(beta_val,2)) * DVA(Ay,Az,C,Bz,d_p,ve);
    double C0_term3 = 2 * alpha_f * DVA(Ay,Az,C,Bz,d_p,ve) * G(W,beta_val,omega_plus,omega_minus,x_f);
    double C0_term4 = 0.5 * (DV(Ay,Az,d_p,ve,C) + DA(Bz,C,ve,d_p) + 2*alpha_f*DVA(Ay,Az,C,Bz,d_p,ve))
                      * (2*H1(beta_val,omega_plus,omega_minus,x_f,W) - H2(beta_val,omega_plus,omega_minus,x_f,W));
    double C0 = (C0_term1+C0_term2) * F(W,beta_val,omega_plus,omega_minus) + C0_term3 + C0_term4;

    double C1_term1 = 2.0 * (2.0*EVA(Ay,ve,Az,Bz,C,d_p) + alpha_f*(1.0-pow(beta_val,2))*EA(Bz,ve,C,d_p)) * F(W,beta_val,omega_plus,omega_minus);
    double C1_term2 = 2.0 * alpha_f * (EV(Ay,Az,C,d_p,ve)+EA(Bz,ve,C,d_p)) * G(W,beta_val,omega_plus,omega_minus,x_f);
    double C1_term3 = -2.0 * (2.0*EVA(Ay,ve,Az,Bz,C,d_p) + alpha_f * (EV(Ay,Az,C,d_p,ve)+EA(Bz,ve,C,d_p)))
                      * H1(beta_val,omega_plus,omega_minus,x_f,W);
    double C1 = C1_term1 + C1_term2 + C1_term3;

    double C2_term1 = 0.5 * ( (3-pow(beta_val,2)) * (DV(Ay,Az,d_p,ve,C)+DA(Bz,C,ve,d_p)) +
                       6*alpha_f*(1-pow(beta_val,2))*DVA(Ay,Az,C,Bz,d_p,ve));
    double C2_term2 = alpha_f*2*DVA(Ay,Az,C,Bz,d_p,ve)*G(W,beta_val,omega_plus,omega_minus,x_f);
    double C2_term3 = -1.5 * (DV(Ay,Az,d_p,ve,C) + DA(Bz,C,ve,d_p) + 2*alpha_f*DVA(Ay,Az,C,Bz,d_p,ve));
    double C2 = C2_term1 * F(W,beta_val,omega_plus,omega_minus) + C2_term2
              + C2_term3 * (2*H1(beta_val,omega_plus,omega_minus,x_f,W) - H2(beta_val,omega_plus,omega_minus,x_f,W));


   
    return (   (C0 + C1*thetaCosine + C2*pow(thetaCosine,2)));

}




void Function(){


TFile *file = TFile::Open("leptons_sm.root");
TTree *ch1 = (TTree*)file->Get("lepton_tree");


Double_t energy, cosTheta;
ch1->SetBranchAddress("energy", &energy);
ch1->SetBranchAddress("cosTheta", &cosTheta);

TH2D *S0 = new TH2D("S0", "F vs xf Costheta; x_f; cos(#theta)",
                           20, 0.1, 1, // x_f range
                           20, -1,1);// cos(theta) range

TH2D *S0_xDim = new TH2D("S0_xDim", "F vs x_f; x_f; ",
                           20, 0.1, 1, // x_f range
                           20, 0,5);


for (Long64_t ientry = 0; ientry < ch1->GetEntries(); ientry++) {
       ch1->GetEntry(ientry);
       double x_f = (2*(energy/mtop)*std::sqrt((1-beta_val)/(1+beta_val)));
       double F = Function_Theta(x_f, cosTheta);

             S0_xDim->Fill(x_f, F);
              S0->Fill(x_f, cosTheta, F);
}

S0_xDim->Draw();

}