#extracts and retrieves the NPs from a workspace. 
#Usage: python extract_np.py path/to/workspace.root

import ROOT
import argparse

def parse_arguments():
    parser = argparse.ArgumentParser(description='Process a ROOT file and retrieve non-constant parameters.')
    parser.add_argument('workspace_file', type=str, help='Path to the workspace file')
    return parser.parse_args()

def main():
    # Parse command line arguments
    args = parse_arguments()

    # Open the chosen ROOT file
    wsFile = ROOT.TFile.Open(args.workspace_file)

    if not wsFile or wsFile.IsZombie():
        print("Error: Unable to open the ROOT file.")
        return

    # Retrieve the workspace
    w = wsFile.Get("w")

    if not w:
        print("Error: Unable to retrieve workspace 'w' from the ROOT file.")
        return

    # Retrieve the ModelConfig from the workspace
    config = w.genobj("ModelConfig")

    if not config:
        print("Error: Unable to retrieve 'ModelConfig' from the workspace.")
        return

    # Get all the PDF variables and parameters of interest (POIs)
    pdfvars = config.GetPdf().getParameters(config.GetObservables())
    pois = config.GetParametersOfInterest()

    # Collect POI names
    pois_names = []
    it_pois = pois.createIterator()
    poi = it_pois.Next()
    while poi:
        pois_names.append(poi.GetName())
        poi = it_pois.Next()

    # Collect non-constant parameters that are not in POIs and are of type RooRealVar
    res = []
    it_vars = pdfvars.createIterator()
    var = it_vars.Next()
    while var:
        if var.GetName() not in pois_names and not var.isConstant() and var.InheritsFrom("RooRealVar"):
            res.append(var.GetName())
        var = it_vars.Next()

    # Output the results
    print("\nFinal list of the {} NPs\n".format(len(res)))
    print(res)
    print("****************************************\n")

if __name__ == "__main__":
    main()