import flavio
import flavio.plots
import flavio.statistics.fits
import matplotlib.pyplot as plt
from collections import OrderedDict
import numpy as np

angobservables, q2bins, Data, label, decay = {}, {}, {}, {}, {}

lhcbangobs = ['<FL>', '<S5>', '<S3>', '<S4>', '<AFB>']
lhcbq2bins = [[0.1, 0.98], [1.1, 2.5], [2.5, 4], [4, 6], [15, 17], [17, 19]]

# previous measurements
angobservables['lhcb15'] = lhcbangobs 
q2bins['lhcb15'] = lhcbq2bins
Data['lhcb15'] = 'LHCb B->K*mumu 2015 S ' #already in flavio's repo
decay['lhcb15'] = '(B0->K*mumu)'
label['lhcb15'] = r'$B^0\to K^{\ast 0}(K^+\pi^-) \mu\mu$ LHCb 3/fb'

with open('./measurements/CMS_2017.yml', 'r') as myfile:
    flavio.measurements._load(myfile.read())
angobservables['cms17'] = ['<P1>', '<P5p>']
q2bins['cms17'] = [[1, 2], [2, 4.3], [4.3, 6], [14.18, 16], [16, 19]]
Data['cms17'] = 'CMS B->K*mumu 2017 P '
decay['cms17'] = '(B0->K*mumu)'
label['cms17'] = r'$B^0\to K^{\ast 0}(K^+\pi^-) \mu\mu$ CMS 20.5/fb'

with open('./measurements/ATLAS_2017.yml', 'r') as myfile:
    flavio.measurements._load(myfile.read())
#angobservables['atlas17'] = ['<P1>', '<P4p>', '<P5p>']
angobservables['atlas17'] = ['<FL>', '<S3>', '<S4>', '<S5>']
q2bins['atlas17'] = [[0.04, 2], [2, 4], [4, 6]]
Data['atlas17'] = 'ATLAS B->K*mumu 2017 S '
decay['atlas17'] = '(B0->K*mumu)'
label['atlas17'] = r'$B^0\to K^{\ast 0}(K^+\pi^-) \mu\mu$ ATLAS 20.3/fb'

# expected sensitivities for B+
with open('./measurements/Bp2Kstmumu_2016.yml', 'r') as myfile:
    flavio.measurements._load(myfile.read())
angobservables['lhcbBp'] = lhcbangobs
q2bins['lhcbBp'] = lhcbq2bins
Data['lhcbBp'] = 'LHCb B+->K*mumu 2016 S '
decay['lhcbBp'] = '(B+->K*mumu)'
label['lhcbBp'] = r'$B^+\to K^{\ast +}(K_S\pi^+) \mu\mu$ LHCb 5.2/fb'

with open('./measurements/Bp2Kstmumu_2018andpi0.yml', 'r') as myfile:
    flavio.measurements._load(myfile.read())
angobservables['lhcbBp18'] = lhcbangobs
q2bins['lhcbBp18'] = lhcbq2bins
Data['lhcbBp18'] = 'LHCb B+->K*mumu 2018 S '
decay['lhcbBp18'] = '(B+->K*mumu)'
label['lhcbBp18'] = r'$B^+\to K^{\ast +}(K_S\pi^+, K^+\pi^0) \mu\mu$ LHCb 9.0/fb'


obss, measurements = {}, {}
for dataset in Data.keys():
    print('now loading dataset '+ dataset)
    measurements[dataset], obss[dataset] = [], []
    for q2 in q2bins[dataset]:
        measurements[dataset].append(Data[dataset]+'{}-{}'.format(q2[0], q2[1]))
        for obs in angobservables[dataset]:
            obss[dataset].append((obs+decay[dataset], q2[0], q2[1]))


def fC9C10(C9, C10):
    return { 'C9_bsmumu': C9, 'C10_bsmumu': C10 }
def fastfit_C9C10(name, obslist, measurements):
    return flavio.statistics.fits.FastFit(
                name = name,
                observables = obslist,
                fit_wc_function = fC9C10,
                input_scale = 4.8,
                include_measurements = measurements,
            )

obs_fastfits_c9c10 = {}
c9lim=[-4,4]
c10lim=[-4,4]

plt.figure(figsize=(6, 6))

plots = {}
for i, f in enumerate(Data.keys()):
    obs_fastfits_c9c10[f] = fastfit_C9C10('C9-C10 fit', obss[f], measurements[f])
    print('making measurement ', f)
    obs_fastfits_c9c10[f].make_measurement(threads=8)
    print('plotting ', f)
    if 'lhcbBp' in f: fill=False
    else: fill=True
    color = i+1
    if color>=5: color = i+2 #avoid yellow
    plots[f] = flavio.plots.likelihood_contour(obs_fastfits_c9c10[f].log_likelihood,
                                    c9lim[0], c9lim[1], c10lim[0], c10lim[1], 
                                    col=color, filled=fill, label=label[f],
                                    interpolation_factor=3, threads=8, steps=20)
    print('done')

plt.ylim(-3, 5)
plt.legend(fontsize=12)
flavio.plots.flavio_branding(x=0.8, y=0.65, version=True)
plt.xlabel(r'Re~$C_9^{\text{NP}}$')
plt.ylabel(r'Re~$C_{10}^{\text{NP}}$')
plt.title(r'$B\to K^\ast\mu^+\mu^-$ angular analyses at the LHC', fontsize=15, pad=15)
plt.plot([0], [0], marker='*', color='black', markersize=10)
sm = plt.text(0.2, -0.3, 'SM', fontsize=12)
plt.savefig('C9_vs_C10.pdf')

plt.ion()
plt.show()