import ROOT


def StandardFrequentistDiscovery(
    infile="",
    workspaceName="channel1",
    modelConfigNameSB="ModelConfig",
    dataName="obsData",
    toys=1000,
    poiValueForBackground=0.0,
    poiValueForSignal=1.0,
):

    # The workspace contains the model for s+b. The b model is "autogenerated"
    # by copying s+b and setting the one parameter of interest to zero.
    # To keep the script simple, multiple parameters of interest or different
    # functional forms of the b model are not supported.

    # for now, assume there is only one parameter of interest, and these are
    # its values:

    # -------------------------------------------------------
    # First part is just to access a user-defined file
    # or create the standard example file if it doesn't exist
    filename = ""
    if infile == "":
        filename = "results/example_channel1_GammaExample_model.root"
        fileExist = not ROOT.gSystem.AccessPathName(filename)  # note opposite return code
        # if file does not exists generate with histfactory
        if not fileExist:
            # Normally this would be run on the command line
            print(f"will run standard hist2workspace example")
            ROOT.gROOT.ProcessLine(".!  prepareHistFactory .")
            ROOT.gROOT.ProcessLine(".!  hist2workspace config/example.xml")
            print(f"\n\n---------------------")
            print(f"Done creating example input")
            print(f"---------------------\n\n")

    else:
        filename = infile

    # Try to open the file
    file = ROOT.TFile.Open(filename)

    # if input file was specified but not found, quit
    if not file:
        print(f"StandardRooStatsDemoMacro: Input file {filename} is not found")
        return -1

    # -------------------------------------------------------
    # Tutorial starts here
    # -------------------------------------------------------

    mn_t = ROOT.TStopwatch()
    mn_t.Start()

    # get the workspace out of the file
    w = file.Get(workspaceName)
    if not w:
        print(f"workspace not found")
        return -1.0

    # get the modelConfig out of the file
    mc = w[modelConfigNameSB]

    # get the data out of the file
    data = w[dataName]

    # make sure ingredients are found
    if not data or not mc:
        w.Print()
        print(f"data or ModelConfig was not found")
        return -1.0

    firstPOI = mc.GetParametersOfInterest().first()
    firstPOI.setVal(poiValueForSignal)
    mc.SetSnapshot(mc.GetParametersOfInterest())
    # create null model
    mcNull = mc.Clone("ModelConfigNull")
    firstPOI.setVal(poiValueForBackground)
    mcNull.SetSnapshot(mcNull.GetParametersOfInterest().snapshot())

    # ----------------------------------------------------
    # Configure a ProfileLikelihoodTestStat and a SimpleLikelihoodRatioTestStat
    # to use simultaneously with ToyMCSampler
    plts = ROOT.RooStats.ProfileLikelihoodTestStat(mc.GetPdf())
    plts.SetOneSidedDiscovery(True)
    plts.SetVarName("q_0/2")

    # ----------------------------------------------------
    # configure the ToyMCImportanceSampler with two test statistics
    toymcs = ROOT.RooStats.ToyMCSampler(plts, 50)

    # Since this tool needs to throw toy MC the PDF needs to be
    # extended or the tool needs to know how many entries in a dataset
    # per pseudo experiment.
    # In the 'number counting form' where the entries in the dataset
    # are counts, and not values of discriminating variables, the
    # datasets typically only have one entry and the PDF is not
    # extended.
    if not mc.GetPdf().canBeExtended():
        if data.numEntries() == 1:
            toymcs.SetNEventsPerToy(1)
        else:
            print(f"Not sure what to do about this model")

    # instantiate the calculator
    freqCalc = ROOT.RooStats.FrequentistCalculator(data, mc, mcNull, toymcs)
    freqCalc.SetToys(toys, toys)  # null toys, alt toys

    # Run the calculator and print result
    freqCalcResult = freqCalc.GetHypoTest()
    freqCalcResult.GetNullDistribution().SetTitle("b only")
    freqCalcResult.GetAltDistribution().SetTitle("s+b")
    freqCalcResult.Print()
    pvalue = freqCalcResult.NullPValue()

    # stop timing
    mn_t.Stop()
    print(f"total CPU time: ", mn_t.CpuTime())
    print(f"total real time: ", mn_t.RealTime())

    # plot
    c1 = ROOT.TCanvas()
    plot = ROOT.RooStats.HypoTestPlot(freqCalcResult, 100, -0.49, 9.51)
    plot.SetLogYaxis(True)
    plot.Draw()
    c1.Update()
    c1.Draw()
    c1.SaveAs("StandardFrequentistDiscovery.1.png")

    # adding chi2 to plot a different plot.
    # plot 1 and plot 2 have problems at HypoTestPlot.AddTF1
    c2 = ROOT.TCanvas("myc2", "myc2")
    nPOI = 1
    # g =  TF1("g", TString.Format("1*ROOT::Math::chisquared_pdf(2*x,{},0)".format( nPOI)), 0, 9) # doesn´t work properly
    scale_ctte = 0.2
    ROOT.gInterpreter.Declare(
        f"""TF1 *g = new TF1("g", TString::Format("{scale_ctte}*ROOT::Math::chisquared_pdf(2*x,%d,0)", {nPOI}), 0, 9);"""
    )
    g = ROOT.g
    g.SetLineColor(ROOT.kBlack)
    g.SetLineStyle(7)
    plot2 = ROOT.RooStats.HypoTestPlot(freqCalcResult, 100, -0.49, 9.51)
    tmptitle = f"#chi^{{2}}(2x,{nPOI})"
    plot2.AddTF1(g, tmptitle, "SAME")
    plot2.SetLogYaxis(True)
    plot2.Draw()

    c2.Update()
    c2.Draw()
    c2.SaveAs("StandardFrequentistDiscovery.2.png")

    return pvalue


StandardFrequentistDiscovery(
    infile="",
    workspaceName="channel1",
    modelConfigNameSB="ModelConfig",
    dataName="obsData",
    toys=1000,
    poiValueForBackground=0.0,
    poiValueForSignal=1.0,
)