import flavio
from wilson import Wilson
from uncertainties import ufloat

angobs = ['FL', 'S3', 'S4', 'S5', 'AFB', 'S7', 'S8', 'S9']
#angobs = ['P1', 'P2', 'P3', 'P4p', 'P5p', 'P6p', 'P8p']

#q2bins = [[0.1, 0.98], [1.1, 2.5], [2.5, 4], [4, 6], [6, 8], [11, 12.5], [15, 17], [17, 19]]
q2bins = [[0.1, 4.0],[4.0,5.0]]#, [4, 8], [11.0,12.5], [15.0,19.0]]
SMpred, NPpred = {}, {} # to be filled

# setting Wilson coefficients for NP
w = Wilson({ 'C9_bsmumu': -1}, scale=160, eft='WET', basis='flavio')

doNP = False
doSM = True
# getting predictions for both SM and NP
for ibin, q2bin in enumerate(q2bins): 
    print(f'Bin {q2bin}') 
    SMpred[ibin], NPpred[ibin] = {}, {}
    if (doSM):
        for o in angobs: 
            SMpred[ibin][o] = ufloat(flavio.sm_prediction(f'<{o}>(B+->K*mumu)',
                                    q2min=q2bin[0], q2max=q2bin[1]),
                                 flavio.sm_uncertainty(f'<{o}>(B+->K*mumu)',
                                    q2min=q2bin[0], q2max=q2bin[1], N=100)) #N=10k for the final plots, N=100 for the fitter
    if (doNP):
        for o in angobs: 
            NPpred[ibin][o] = ufloat(flavio.np_prediction(f'<{o}>(B+->K*mumu)', w, 
                                    q2min=q2bin[0], q2max=q2bin[1]),
                                 flavio.np_uncertainty(f'<{o}>(B+->K*mumu)', w, 
                                    q2min=q2bin[0], q2max=q2bin[1]))

# Just printing everything out in human-readable format here
def printout(dic):
    for ibin, q2bin in enumerate(q2bins): 
        print(f'Bin {q2bin}') 
        for o in angobs: 
            if o=='FL':
                print(f'\tS1s = {3/4*(1-dic[ibin][o])}')
            elif o=='AFB':
                print(f'\tS6s = {4/3*dic[ibin][o]}')
            else:
                print(f'\t{o} = {dic[ibin][o]}')

def printoutCpp(dic):
    for o in angobs:
        if o=='FL':
            print('\tS1s[8] = {')
        elif o=='AFB':
            print('\tS6s[8] = {')
        else:
            print(f'\t{o}[8] = {{')
        for ibin, q2bin in enumerate(q2bins):
            if o=='FL':
                print(f'{3/4*(1-dic[ibin][o].n)},', end='')
            elif o=='AFB':
                print(f'{4/3*dic[ibin][o].n},', end='')
            else:
                print(f'{dic[ibin][o].n},', end='')
        print(f'}};')
    for o in angobs:
        if o=='FL':
            print('\tS1s_error[8] = {')
        elif o=='AFB':
            print('\tS6s_error[8] = {')
        else:
            print(f'\t{o}_error[8] = {{')
        for ibin, q2bin in enumerate(q2bins):
            if o=='FL':
                print(f'{3/4*(dic[ibin][o].s)},', end='')
            elif o=='AFB':
                print(f'{4/3*dic[ibin][o].s},', end='')
            else:
                print(f'{dic[ibin][o].s},', end='')
        print('};')


if(doSM):
    print('############## SM predictions #################')
    printout(SMpred)

    printoutCpp(SMpred)

    print('\n\n\n')

if(doNP):
    print('############## NP predictions #################')
    printout(NPpred)

    printoutCpp(NPpred)