%matplotlib inline
from __future__ import division
from matplotlib import pyplot as pl
import numpy as np

wWidth = 2.085
zWidth = 2.5

wMass = np.random.normal(80.39, wWidth, 1000000)
zMass = np.random.normal(91.2, zWidth, 1000000)
masses = np.concatenate((wMass, zMass))

wMin, wMax = np.percentile(wMass, 5), np.percentile(wMass, 95)
cut = wMass[np.logical_and(wMass>wMin, wMass < wMax)]
bins = np.linspace(min(wMass), max(wMass), 100)
pl.hist(cut, bins=bins, histtype='step', linewidth=1.5)
pl.hist(wMass, bins=bins, histtype='step', linewidth=1.5)
pl.show()

print "Fraction of width:", np.std(cut)/np.std(wMass)

# let's generate two distributions with a resolution 2 GeV wider
wMassWide1 = np.random.normal(80.39, wWidth+2, 1000000)
zMassWide1 = np.random.normal(91.2, zWidth+2, 1000000)
masses1 = np.concatenate((wMassWide1, zMassWide1))

# now let's generate two distributions with a twice as wide resolution
wMassWide2 = np.random.normal(80.39, wWidth*2, 1000000)
zMassWide2 = np.random.normal(91.2, zWidth*2, 1000000)
masses2 = np.concatenate((wMassWide2, zMassWide2))

# and now we correct for the fact that RMS90 is only <79% of a Gaussian width
wMassWide3 = np.random.normal(80.39, (wWidth*2)/.79, 1000000)
zMassWide3 = np.random.normal(91.2, (zWidth*2)/.79, 1000000)
masses3 = np.concatenate((wMassWide3, zMassWide3))

bins = np.linspace(min(wMassWide3), max(zMassWide3), 200)
pl.hist(masses, bins=bins, histtype='step', label='natural widths', linewidth=2.0, color='k')
pl.hist(masses1, bins=bins, histtype='step', label='2 GeV + natural widths', linewidth=2.0, color='b')
pl.hist(masses2, bins=bins, histtype='step', label='2x natural widths', linewidth=2.0, color='g')
pl.hist(masses3, bins=bins, histtype='step', label='RMS90 = 2x natural widths', linewidth=2.0, color='r')
pl.xlabel('mass (GeV)')
pl.ylabel('arbitrary units')
pl.tick_params(axis='y', labelleft='off')
pl.legend()

HALFWAYPOINT = (91.2-80.39)/2 + 80.39
print("Half way:", HALFWAYPOINT)
print("Wrong assignment: W:", np.sum(wMass>HALFWAYPOINT)/len(wMass), '\tZ:', np.sum(zMass<HALFWAYPOINT)/len(zMass))

print("Wrong assignment: W:", np.sum(wMassWide1>HALFWAYPOINT)/len(wMassWide1), '\tZ:', np.sum(zMassWide1<HALFWAYPOINT)/len(zMassWide1))

print("Wrong assignment: W:", np.sum(wMassWide2>HALFWAYPOINT)/len(wMassWide2), '\tZ:', np.sum(zMassWide2<HALFWAYPOINT)/len(zMassWide2))

print("Wrong assignment: W:", np.sum(wMassWide3>HALFWAYPOINT)/len(wMassWide3), '\tZ:', np.sum(zMassWide3<HALFWAYPOINT)/len(zMassWide3))