import ROOT
#from ROOT import TTree, TFile, Double
import numpy as np
from pdg_const import pdg
import matplotlib
matplotlib.use("Qt5Agg")
import matplotlib.pyplot as plt
        

class model:
    
    def __init__(self):
    

        self.mmu = pdg['muon_M']
        self.mb = pdg["bquark_M"]
        self.ms = pdg["squark_M"]
        self.mK = pdg["Ks_M"]
        self.mB = pdg["Bplus_M"]

        self.C7eff = pdg["C7eff"]
        self.C9eff = pdg["C9eff"]
        self.C10eff = pdg["C10eff"]
        
        #self.C1 = pdg["C1"]
        #self.C2 = pdg["C2"]
        #self.C3 = pdg["C3"]
        #self.C4 = pdg["C4"]
        
        self.GF = pdg["GF"] #Fermi coupling const.
        self.alpha_ew = pdg["alpha_ew"]
        self.Vts = pdg["Vts"]
        self.Vtb = pdg["Vtb"]
        
    def formfactor(self, subscript):
        
        # x = q2
        
        ffplus = ""
        
        if subscript != "0" and subscript != "+" and subscript != "T":
            raise ValueError('Wrong subscript entered, choose either 0, + or T')
        
        #get constants
        
        mh = self.mK
        mbstar0 = pdg["mbstar0"]
        mbstar = pdg["mbstar"]
        b0 = pdg["b0"]
        bplus = pdg["bplus"]
        bT = pdg["bT"]
        mB = self.mB
        mK = self.mK
        
        N = 3
        
        #some helperfunctions
        
        tpos = "(({0} - {1})*({0} - {1}))".format(mB, mK)
        tzero = "(({0} + {1})*(sqrt({0})-sqrt({1}))^2)".format(mB, mK)
        
        z_oben = "(sqrt({0} - pow(x,2)) - sqrt({0} - {1}))".format(tpos, tzero)
        z_unten = "(sqrt({0} - pow(x,2)) + sqrt({0} - {1}))".format(tpos, tzero)
        z = "({0})/({1})".format(z_oben, z_unten)
        
        #calculate f0
        
        if subscript == "0":
        
            prefactor_ff0 = "1/(1 - pow(x,2)/({0}^2))".format(mbstar0)
            ff0 = ""
            
            for i in range(N):
                ff0 += "{0}*({1}^{2})".format(b0[i], z, i)
                if i < (N -1):
                    ff0 += "+"
            
            ff0 = prefactor_ff0 + "*" + ff0
            
            return "(" + ff0  + ")"
        
        
        #calculate f+
        
        elif subscript == "+":
        
            ffplus = ""
            prefactor_ffplus = "1/(1 - pow(x,2)/(pow({0},2)))".format(mbstar)
            
            for i in range(N):
                ffplus += "{0} * (pow({1},{2}) - (pow((-1),({2}-3)) * ({2}/3) * pow({1},3)))".format(bplus[i], z, i)
                if i < (N -1):
                    ffplus += "+"
            
            ffplus = prefactor_ffplus + "*" + ffplus
            
            return "(" + ffplus + ")"
        
        #calculate fT
        
        else:
        
            ffT = ""
            prefactor_ffT = "1/(1 - pow(x,2)/(pow({0},2)))".format(mbstar)
                        
            for i in range(N):
                ffplus += "{0} * (pow({1},{2}) - (pow((-1),({2}-3)) * ({2}/3) * pow({1},3)))".format(bT[i], z, i)
                if i < (N -1):
                    ffT += "+"
            
            ffT = prefactor_ffT + "*" + ffT
            print(ffT)
            
            return "(" + ffT + ")"
        
        
        
    def axiv_nonres(self):
        
        # [0] = amp
        
        GF = self.GF
        alpha_ew = self.alpha_ew
        Vtb = self.Vtb
        Vts = self.Vts
        C10eff = self.C10eff
        
        mmu = self.mmu
        mb = self.mb
        ms = self.ms
        mK = self.mK
        mB = self.mB        
    
        #Some helperfunctions
        
        beta = "sqrt(1. - 4. * pow({0},2) / x))".format(mmu)
        
        kabs = "(sqrt(pow({0},2) + pow(x,2)/pow({0},2) + pow({1},4)/pow({0},2) - 2 * (pow({0},2) * pow({1},2) + pow({1},2) * x + pow({0},2) * x) / pow({0},2)))".format(mB, mK)
        
        #prefactor in front of whole bracket
        
        prefactor1 = "(pow({0},2) * pow({1},2) * pow({2}*{3},2) * {4} * {5} / (128. * pow({6},5)))".format(GF, alpha_ew, Vtb, Vts, kabs, beta, np.pi)
        
        #left term in bracket
        
        bracket_left = "(2./3. * pow({0},2) * pow({1},2) * pow({2} * ({3}),2))".format(kabs, beta, C10eff, self.formfactor("+"))
        
        #middle term in bracket
        
        _top = "(4. * pow({0},2) * pow((pow({1},2) - pow({2},2)),2))".format(mmu, mB, mK)
        
        _under = "(x * pow({0},2))".format(mB)
        
        bracket_middle = "({0}/{1} * (sqrt((({2} * {3}) * ({2} * {3}))) *sqrt((({2} * {3}) * ({2} * {3})))))".format(_top, _under, C10eff, self.formfactor("0"))
        
        axiv = prefactor1 + "*" + "(" + bracket_left + "+" + bracket_middle + ") * 2 * sqrt(x) * [0]"
        
        return "(" + axiv + ")"
    

    def vec_nonres(self):
        
        # [1] = amp
        
        GF = self.GF
        alpha_ew = self.alpha_ew
        Vtb = self.Vtb
        Vts = self.Vts
        C7eff = self.C7eff
        C9eff = self.C9eff
        
        mmu = self.mmu
        mb = self.mb
        ms = self.ms
        mK = self.mK
        mB = self.mB        
    
        #Some helperfunctions
        
        beta = "sqrt(1. - 4. * pow({0},2) / x))".format(mmu)
        
        kabs = "(sqrt(pow({0},2) + pow(x,2)/pow({0},2) + pow({1},4)/pow({0},2) - 2 * (pow({0},2) * pow({1},2) + pow({1},2) * x + pow({0},2) * x) / pow({0},2)))".format(mB, mK)
        
        #prefactor in front of whole bracket
        
        prefactor1 = "(pow({0},2) * pow({1},2) * pow({2}*{3},2) * {4} * {5} / (128. * pow({6},5)))".format(GF, alpha_ew, Vtb, Vts, kabs, beta, np.pi)
        
        #right term in bracket
        
        prefactor2 = "(pow({0},2) * (1. - 1./3. * pow({1},2)))".format(kabs, beta)
        
        abs_bracket = "pow(abs({0} * {1} + 2 * {2} * ({3} + {4})/({5} + {6}) * {7}),2)".format(C9eff, self.formfactor("+"), C7eff, mb, ms, mB, mK, self.formfactor("T"))
                
        bracket_right = "{0} * {1}".format(prefactor2, abs_bracket)
        
        vec =  "[1] * {0} * {1} * 2 * sqrt(x)".format(prefactor1, bracket_right)
        
        return "(" + vec + ")"
    
    def total_nonres(self):
        
        #[0] = amp axiv
        #[1] = vec amp
        #[2] = tot nonres amp
        
        #Get constants
        
        GF = self.GF
        alpha_ew = self.alpha_ew
        Vtb = self.Vtb
        Vts = self.Vts
        C10eff = self.C10eff
        C9eff = self.C9eff
        C7eff = self.C7eff
        
        mmu = self.mmu
        mb = self.mb
        ms = self.ms
        mK = self.mK
        mB = self.mB
        
        #vector nonresonant part
        
        vec_nonres_part = self.vec_nonres()
        
        #axial verctor nonresonant part including C7
        
        axiv_nonres_part = self.axiv_nonres()
        
        #Complete term
        
        return "[2] * ({0} + {1})".format(vec_nonres_part, axiv_nonres_part)
    

modl = model()

mB = modl.mB
mK = modl.mK
mmu = modl.mmu

tpos = "(({0} - {1})*({0} - {1}))".format(mB, mK)
tzero = "(({0} + {1})*(sqrt({0})-sqrt({1}))^2)".format(mB, mK)

z_oben = "(sqrt({0} - x) - sqrt({0} - {1}))".format(tpos, tzero)
z_unten = "(sqrt({0} - x) + sqrt({0} - {1}))".format(tpos, tzero)
z = "({0})/({1})".format(z_oben, z_unten)

#print(z_oben)

#z_oben = "sqrt(" +  tpos + "-x) - sqrt(" + tpos + "-" + tzero + ")"
#print(z_oben)




dq = np.linspace(2*mmu, (mB-mK) - 0.1 ,1000)

dq2 = []

for i in dq:
	dq2.append(i**2)
	
vec_fit = ROOT.TF1("vec_fit", modl.formfactor("T"), 0, (mB-mK) - 0.1)

vec_fit.SetParameter(0,1)
vec_fit.SetParameter(1,1)
vec_fit.SetParameter(2,1)

canvas = ROOT.TCanvas("canvas")

canvas.cd()

vec_fit.Draw()

canvas.Print("root_test.pdf")

canvas.Close()


