import awkward
import numpy as np
from coffea import nanoevents
events = nanoevents\
         .NanoEventsFactory\
         .from_root('../tests/samples/nano_dy.root')\
         .events()

def some_event_weight(ones):
    return (1.0 + np.array([0.05, -0.05], dtype=np.float32)) * ones[:, None]

events.add_systematic("RenFactScale", "UpDownSystematic", "weight", some_event_weight)
events.add_systematic("XSectionUncertainty", "UpDownSystematic", "weight", some_event_weight)

muons = events.Muon
jets = events.Jet

def muon_pt_scale(pt):   
    return (1.0 + np.array([0.05, -0.05], dtype=np.float32)) * pt[:, None]

def muon_pt_resolution(pt):   
    return np.random.normal(pt[:,None], np.array([0.02, 0.01], dtype=np.float32))

def muon_eff_weight(ones):
    return (1.0 + np.array([0.05, -0.05], dtype=np.float32)) * ones[:, None]

muons.add_systematic("PtScale", "UpDownSystematic", "pt", muon_pt_scale)
muons.add_systematic("PtResolution", "UpDownSystematic", "pt", muon_pt_resolution)
muons.add_systematic("EfficiencySF", "UpDownSystematic", "weight", muon_eff_weight)

def jet_pt_scale(pt):   
    return (1.0 + np.array([0.10, -0.10], dtype=np.float32)) * pt[:, None]

def jet_pt_resolution(pt):   
    return np.random.normal(pt[:,None], np.array([0.20, 0.10], dtype=np.float32))

jets.add_systematic("PtScale", "UpDownSystematic", "pt", jet_pt_scale)
jets.add_systematic("PtResolution", "UpDownSystematic", "pt", jet_pt_resolution)

awkward.flatten(muons.pt)

awkward.flatten(muons.systematics.PtScale.up)

awkward.flatten(muons.systematics.PtScale.up.pt)

awkward.flatten(jets.pt)

awkward.flatten(jets.systematics.PtScale.up)

awkward.flatten(jets.systematics.PtScale.up.pt)

events.systematics.RenFactScale.up.weight_RenFactScale

# TODO: Make it so that syst_muons.Y > X returns boolean values
#       for all variations over Y. 
#       Requires some tracking of (pieces of) "what".