import ROOT


def rs301_splot():

    lowRange = 5600
    highRange = 5900
    
    x = ROOT.RooRealVar("lab0_M", "lab0 MASS", lowRange, highRange, "MeV")
    isolation = ROOT.RooRealVar("lab0_IPCHI2_OWNPV", "isolation", 0., 100, "MeV")

 #   filename = ROOT.TFile("BDT_Xib0.root")
   # tree = filename.Get("mytree")
    hist = ROOT.TH1F("hist","hist title",100,lowRange,highRange)
    hist.FillRandom("gaus")
    data = ROOT.RooDataHist("data","my dataset",x,hist)
  #  tree.Project("hist","lab0_M")
   

    print ("make SIGNAL model")

    mean = ROOT.RooRealVar("mean", "Mass mean", 5800, 5750, 5830)
    width = ROOT.RooRealVar("width","gaussian width", 25,10,100)
    gauss = ROOT.RooGaussian("gauss","gaussian pdf",x,mean,width)

    print ("make BACKGROUND model")

    c0 = ROOT.RooRealVar("c0","coeff #0",0.8,-1.,1.)
    c1 = ROOT.RooRealVar("c1","coeff #1",0.1,-1.,1.)
    c2 = ROOT.RooRealVar("c2","coeff #2",-0.1,-1.,1.)
    ground=ROOT.RooChebychev("bkg","background p.d.f.",x,ROOT.RooArgList(c0,c1,c2)) 

    print("make FULL model")
    
    sYield = ROOT.RooRealVar("sYield", "fitted signal yield", 10000,0.,500000.)
    bYield = ROOT.RooRealVar("bYield", "fitted bkg yield", 80000,0.,500000.)

    # now make the combined model
    model = ROOT.RooAddPdf("model","MODEL", ROOT.RooArgList(gauss,ground),ROOT.RooArgList(sYield,bYield))

    # make the toy data
    print ("make data set and import to workspace")



    print ("Calculate sWeights")


    # fit the model to the data.
    model.fitTo(data, ROOT.RooFit.Extended())


    sData = ROOT.RooStats.SPlot("sData", "An SPlot", data, model, ROOT.RooArgList(sYield, bYield))

    # Check that our weights have the desired properties

    print ("Check SWeights:")
    print("")
    print ("make plots")

    # make our canvas
    cdata = ROOT.TCanvas("sPlot", "sPlot demo", 400, 600)
    cdata.Divide(1, 3)

    model.fitTo(data, ROOT.RooFit.Extended())

    cdata = ROOT.TCanvas()
    cdata.cd(1)
    frame = x.frame()
    data.plotOn(frame)
    model.plotOn(frame)
    ras_gauss = ROOT.RooArgSet(gauss)
    ras_ground = ROOT.RooArgSet(ground)
    model.plotOn(frame, ROOT.RooFit.Components(
        ras_gauss), ROOT.RooFit.LineStyle(ROOT.kDashed), ROOT.RooFit.LineColor(ROOT.kRed))
    model.plotOn(frame, ROOT.RooFit.Components(
        ras_ground), ROOT.RooFit.LineStyle(ROOT.kDashed), ROOT.RooFit.LineColor(ROOT.kGreen))

    frame.SetTitle("Fit of model to discriminating variable")
    frame.Draw()

   
    cdata.cd(2)

    dataw_s = ROOT.RooDataSet(data.GetName(), data.GetTitle(),
                           data, data.get(), "", "sigYield_sw")

    frame2 = isolation.frame()
    dataw_s.plotOn(frame2, ROOT.RooFit.DataError(ROOT.RooAbsData.SumW2))

    frame2.SetTitle("isolation distribution for S")
    frame2.Draw()

    cdata.cd(3)
    dataw_b = ROOT.RooDataSet(data.GetName(), data.GetTitle(), data, data.get(), "", "bkgYield_sw")
    frame3 = isolation.frame()
    dataw_b.plotOn(frame3, ROOT.RooFit.DataError(ROOT.RooAbsData.SumW2))

    frame3.SetTitle("isolation distribution for B")
    frame3.Draw()

    cdata.SaveAs("rs301_splot.pdf")


if __name__ == "__main__":
    rs301_splot()