import ROOT
import aa

#declare function to load atmospheric flux
ROOT.gInterpreter.ProcessLine("#include \"Flux.hh\"")
func_def="""
map<string,int> nupdgs  = { {\"nue\",12},{\"numu\",14},{\"antinue\",-12},{\"antinumu\",-14} };

Flux_Honda2006 conv;
double AanetConv(int pdg, double e, double cth) { 
    return conv.dNdEdOmega(pdg,TMath::Log10(e),cth); 
}
Flux_prompt_Sarsevic_std prompt;
double AanetPrompt(int pdg, double e, double cth) { 
    return prompt.dNdEdOmega(pdg,TMath::Log10(e),cth); 
}
Flux_Atmospheric atm;
double AanetAtm(int pdg, double e, double cth) { 
    return atm.dNdEdOmega(pdg,TMath::Log10(e),cth); 
}

double angle(double dx1, double dy1, double dz1, double dx2, double dy2, double dz2) {
  double n1 = TMath::Sqrt(TMath::Power(dx1,2)+TMath::Power(dy1,2)+TMath::Power(dz1,2) );
  double n2 = TMath::Sqrt(TMath::Power(dx2,2)+TMath::Power(dy2,2)+TMath::Power(dz2,2) );
  double dot = (dx1*dx2+dy1*dy2+dz1*dz2)/n1/n2;
  return TMath::RadToDeg()*TMath::ACos(dot);
}

double myratio(double num, double den) {
  if (den==0) return 0.;
  else        return num/den;
}

"""
ROOT.gInterpreter.Declare(func_def)



#########################################################################################
#########################################################################################

class Flux :

  def StringFlux(self,pdg,e,cth) :
    return ""

  def ComputeFlux(self,pdg,e,cth) :
      return 1.

  def Rate(self,pdg,haeff) :
    Nnus = 0
    for ie in range( 1, haeff.GetNbinsX()+1 ) :  
        Ngevs = ROOT.TMath.Power(10,haeff.GetXaxis().GetBinUpEdge(ie)) - ROOT.TMath.Power(10,haeff.GetXaxis().GetBinLowEdge(ie))
        e = ROOT.TMath.Power(10.,haeff.GetXaxis().GetBinCenter(ie))
        for ic in range( 1, haeff.GetNbinsY()+1 ) : 
          cth = haeff.GetYaxis().GetBinCenter(ic)
          Nnus += haeff.GetBinContent(ie, ic) * self.ComputeFlux(pdg,e,cth) * Ngevs
    return Nnus * 4.*ROOT.TMath.Pi()/haeff.GetNbinsY()


#########################################################################################
#########################################################################################

class IsotropicFlux(Flux) :

  def __init__(self,norm,gamma,cutoff) :
    self.norm   = norm
    self.gamma  = gamma
    self.cutoff = cutoff

  def StringFlux(self,pdg,e,cth) :
    return str(self.norm)+"*TMath::Power("+e+",-"+str(self.gamma)+")*TMath::Exp(-"+e+"/"+str(self.cutoff)+")"

  def ComputeFlux(self,pdg,e,cth) :
      return self.norm * ROOT.TMath.Power(e,-self.gamma) * ROOT.TMath.Exp(-e/self.cutoff)


#########################################################################################
#########################################################################################

class AanetFlux(Flux) :

  def __init__(self,fname) :
    self.fname   = fname

  def StringFlux(self,pdg,e,cth) :
    return self.fname+"("+pdg+","+e+","+cth+")"

  def ComputeFlux(self,pdg,e,cth) :
    if   self.fname=="AanetConv"   : return ROOT.AanetConv(pdg,e,cth)
    elif self.fname=="AanetPrompt" : return ROOT.AanetPrompt(pdg,e,cth)
    elif self.fname=="AanetAtm"    : return ROOT.AanetAtm(pdg,e,cth)