import ROOT

def makeMomentMorphing(ws, width_A, width_B, width_i):
        label=""

        m_yy=ws.var("m_yy")

        pdf_list = ROOT.RooArgList()
        pdf_A = ws.function("int_yy_1")
        pdf_B = ws.function("int_yy_2")
        pdf_i = ws.function("int_yy_i")
        #pdf_A = ws.pdf("int_yy_1")
        #pdf_B = ws.pdf("int_yy_2")
        #pdf_i = ws.pdf("int_yy_i")
	pdf_list.add(pdf_A)
	pdf_list.add(pdf_B)

        par_name = "lambda"
        par_title = "#lambda"

        Lambda_list = [width_A, width_B]

        Lambda_0 = Lambda_list[0]
        Lambda_1 = Lambda_list[-1]
	Lambda_i = width_i

        Lambda = ROOT.RooRealVar(par_name, par_title, Lambda_i, Lambda_0, Lambda_1)

        paramVec = ROOT.TVectorD(len(Lambda_list))
        i=0
	for l in Lambda_list:
                paramVec[i] = l
                i=i+1

        #pdf_morph_2 = ROOT.RooMomentMorph('pdf_morph'+label,'pdf_morph'+label,Lambda, ROOT.RooArgList(m_yy), pdf_list, paramVec)
        pdf_morph_2 = ROOT.RooMomentMorphFunc('pdf_morph'+label,'pdf_morph'+label,Lambda, ROOT.RooArgList(m_yy), pdf_list, paramVec)

        c1 = ROOT.TCanvas()
        c1.cd()

        frame_yy = m_yy.frame()
        frame_yy.GetYaxis().SetTitle("Normalized to 1")
        pdf_A.plotOn(frame_yy,ROOT.RooFit.LineColor(ROOT.kBlue), ROOT.RooFit.Name("pdf_A"));
        pdf_B.plotOn(frame_yy,ROOT.RooFit.LineColor(ROOT.kRed), ROOT.RooFit.Name("pdf_B"));
        pdf_i.plotOn(frame_yy,ROOT.RooFit.LineColor(ROOT.kGreen), ROOT.RooFit.Name("pdf_i"));
        pdf_morph_2.plotOn(frame_yy,ROOT.RooFit.LineColor(ROOT.kBlack), ROOT.RooFit.Name("pdf_morph_i"));

	'''
        n = 5
	for i in range(1,n):
                d = (Lambda_1 - Lambda_0)/n
                Lambda.setVal(i*d)
                print(Lambda.getVal())
                pdf_morph_2.plotOn(frame_yy,ROOT.RooFit.LineColor(ROOT.kBlack), ROOT.RooFit.LineStyle(ROOT.kDashed), ROOT.RooFit.Name("pdf_morph_"+str(i)));

        Lambda.setVal(Lambda_i)
        pdf_morph_2.plotOn(frame_yy,ROOT.RooFit.LineColor(ROOT.kBlack), ROOT.RooFit.Name("pdf_morph_i"));

        Lambda.setVal(Lambda_0)
        pdf_morph_2.plotOn(frame_yy,ROOT.RooFit.LineColor(ROOT.kCyan), ROOT.RooFit.LineStyle(ROOT.kDashed), ROOT.RooFit.Name("pdf_morph_A"));
        Lambda.setVal(Lambda_1)
        pdf_morph_2.plotOn(frame_yy,ROOT.RooFit.LineColor(ROOT.kMagenta), ROOT.RooFit.LineStyle(ROOT.kDashed), ROOT.RooFit.Name("pdf_morph_B"));
	'''

        frame_yy.Draw()
	c1.SaveAs("plot_noOffset.png")



w = ROOT.RooWorkspace("ws")

#w.factory("CEXPR::int_yy_1('100-N*TMath::Exp(-b1/1000.*m_yy)*(m_yy*m_yy-mH*mH) / ((m_yy*m_yy-mH*mH)*(m_yy*m_yy-mH*mH) + Gamma_1*Gamma_1*mH*mH)', {m_yy[125,110,140], N[20000], b1[0], mH[125], Gamma_1[1]})")
#w.factory("CEXPR::int_yy_2('100-N*TMath::Exp(-b1/1000.*m_yy)*(m_yy*m_yy-mH*mH) / ((m_yy*m_yy-mH*mH)*(m_yy*m_yy-mH*mH) + Gamma_2*Gamma_2*mH*mH)', {m_yy, N, b1, mH, Gamma_2[10]})")
#w.factory("CEXPR::int_yy_i('100-N*TMath::Exp(-b1/1000.*m_yy)*(m_yy*m_yy-mH*mH) / ((m_yy*m_yy-mH*mH)*(m_yy*m_yy-mH*mH) + Gamma_i*Gamma_i*mH*mH)', {m_yy, N, b1, mH, Gamma_i[5]})")

w.factory("cexpr::int_yy_1('-N*TMath::Exp(-b1/1000.*m_yy)*(m_yy*m_yy-mH*mH) / ((m_yy*m_yy-mH*mH)*(m_yy*m_yy-mH*mH) + Gamma_1*Gamma_1*mH*mH)', {m_yy[125,110,140], N[20000], b1[0], mH[125], Gamma_1[1]})")
w.factory("cexpr::int_yy_2('-N*TMath::Exp(-b1/1000.*m_yy)*(m_yy*m_yy-mH*mH) / ((m_yy*m_yy-mH*mH)*(m_yy*m_yy-mH*mH) + Gamma_2*Gamma_2*mH*mH)', {m_yy, N, b1, mH, Gamma_2[10]})")
w.factory("cexpr::int_yy_i('-N*TMath::Exp(-b1/1000.*m_yy)*(m_yy*m_yy-mH*mH) / ((m_yy*m_yy-mH*mH)*(m_yy*m_yy-mH*mH) + Gamma_i*Gamma_i*mH*mH)', {m_yy, N, b1, mH, Gamma_i[5]})")

makeMomentMorphing(w, 1, 10, 5)