workspace

#!/usr/bin/env python

import sys
from mva import cmd
from mva.samples import Higgs
from mva.workspace import write_workspaces, cuts_workspace, mva_workspace
from mva.analysis import get_analysis
from mva.massregions import MassRegions

parser = cmd.analysis_parser(cmd.general_parser(
    multi_years=True, multi_categories=True))
parser.add_argument('type', choices=('mva', 'cuts'), default='mva')
parser.add_argument('--clf-mass', type=int, default=None)
parser.add_argument('--sideband', default=False, action='store_true')
args = parser.parse_args()

if args.masses == 'all':
    args.masses = Higgs.MASSES
else:
    args.masses = map(int, args.masses.split(','))

params = {}
if args.type == 'mva':
    workspace_func = mva_workspace
    params['clf_mass'] = args.clf_mass
else:
    workspace_func = cuts_workspace

years = args.years
categories = args.categories
if len(categories) == 1 and len(years) > 1:
    categories = categories * len(years)
elif len(categories) != len(years):
    sys.exit("specify the same number of category definitions as years")

cuts = None
if args.sideband:
    massregions = MassRegions()
    cuts = massregions.control_region
    # always unblind in the mass sideband
    args.unblind = True

signal_regions = {}
control_regions = {}
for year, year_categories in zip(years, categories):
    analysis = get_analysis(args, year=year)
    sr, cr = workspace_func(analysis, year_categories, args.masses,
                            unblind=args.unblind,
                            systematics=args.systematics,
                            cuts=cuts,
                            **params)
    signal_regions[year] = sr
    control_regions[year] = cr

suffix = analysis.get_suffix(year=False)
suffix += '_' + args.type
if args.sideband:
    suffix += '_sideband'
if args.output_suffix:
    suffix += '_' + args.output_suffix
path = 'workspaces/hh{0}'.format(suffix.lower())
write_workspaces(path, 'hh', signal_regions, controls=control_regions)