import re
import os
import zfit
import pprint
import zutils.utils as zut
import utils_noroot as utnr
import hqm.model    as hqm_model

from importlib.resources import files
from rkex_model          import model
from hqm.tools.scales    import scales                   as scl
from stats.average       import average                  as stav
from logzero             import logger                   as log
from builder             import builder                  as cb_builder
from misID_tools.zmodel  import misID_real_model_builder as msid

#---------------------------------------------------------------
class rk_model(model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self._d_val = {} 
        self._d_var = {} 

        self._norm_vers = 'v1'
    #---------------------------------------------------------------
    def _get_dset_trig(self, preffix):
        [chan, dset, kind, _ ] = preffix.split('_')
    
        if   chan == 'mm':
            trig = 'MTOS'
        elif chan == 'ee' and kind == 'TOS': 
            trig = 'ETOS'
        elif chan == 'ee' and kind == 'TIS': 
            trig = 'GTIS'
        else:
            log.error(f'Invalid channel: {chan}')
            raise ValueError
    
        return dset, trig
    #---------------------------------------------------------------
    def _is_good_constraint(self, name):
        '''
        TIS muon is not used (PDF dropped in base class), no need for constraints for these
        parameters
        '''
        if '_mm_' in name and '_TIS_' in name:
            return False

        return True
    #---------------------------------------------------------------
    def _preffix_to_info(self, preffix):
        regex = '(ee|mm)_(r1|r2p1|2017|2018|all)_(TOS|TIS)_.*'
        mtch  = re.match(regex, preffix)
        try:
            [chan, dset, trg] = mtch.groups()
        except:
            log.error(f'Cannot extract dataset and trigger from: {preffix}')
            raise

        trig = {'mm_TOS' : 'MTOS', 'mm_TIS' : 'MTOS', 'ee_TOS' : 'ETOS', 'ee_TIS' : 'GTIS'} [f'{chan}_{trg}']

        return dset, trig
    #---------------------------------------------------------------
    def _add_constraints(self, d_cns):
        d_val = { name : cns[0]    for name, cns in d_cns.items() if self._is_good_constraint(name)}
        d_var = { name : cns[1]**2 for name, cns in d_cns.items() if self._is_good_constraint(name)}

        self._d_val.update(d_val)
        self._d_var.update(d_var)
    #---------------------------------------------------------------
    def _get_combinatorial(self, preffix, nent):
        dset, trig    = self._get_dset_trig(preffix) 
        obj           = cb_builder(dset=dset, trigger=trig, vers='v5', q2bin='high', const=False)
        obj.cache_path= './cb_buider.tar.gz'
        cmb, d_cns    = obj.get_pdf(obs=self._obs, preffix=f'cb_{preffix}') 

        self._add_constraints(d_cns)

        name  = r'$\mu\mu$' if preffix.startswith('mm_') else 'ee'
        nent  = self._get_entries(nsig=None, kind='cmb', trigger=trig, year=dset)
        ncb   = zfit.Parameter(f'ncb_{preffix}', nent, 0, 100000)
        cmb.set_yield(ncb)

        return cmb 
    #---------------------------------------------------------------
    def _get_rare_scale(self, year, trigger, kind):
        if   year == 'r1':
            obj_1        = scl(dset='2011', trig=trigger, kind=kind)
            val_1, err_1 = obj_1.get_scale()

            obj_2        = scl(dset='2012', trig=trigger, kind=kind)
            val_2, err_2 = obj_2.get_scale()
        elif year == 'r2p1':
            obj_1        = scl(dset='2015', trig=trigger, kind=kind)
            val_1, err_1 = obj_1.get_scale()

            obj_2        = scl(dset='2016', trig=trigger, kind=kind)
            val_2, err_2 = obj_2.get_scale()
        else:
            obj          = scl(dset=  year, trig=trigger, kind=kind)
            val, err     = obj.get_scale()

            return val, err

        val, err, _ = stav([val_1, val_2], [err_1, err_2])

        return val, err
    #---------------------------------------------------------------
    def _get_entries_from_json(self, kind=None, trig=None, year=None):
        json_path = files('extractor_data').joinpath(f'{self._norm_vers}/{year}_{trig}.json')
        if not os.path.isfile(json_path):
            log.error(f'Cannot find: {json_path}')
            raise FileNotFoundError

        d_par = utnr.load_json(json_path)
        for key, [val, _] in d_par.items():
            flg_1 = key.startswith('ncb_') and kind == 'cmb'
            flg_2 = key.startswith('npr_') and kind == 'prc'
            if flg_1 or flg_2:
                return val

        log.error(f'Cannot find yield for {kind} among:')
        pprint.pprint(d_par)
        raise
    #---------------------------------------------------------------
    def _get_entries(self, nsig=None, kind=None, trigger=None, year=None):
        if   kind in ['prc', 'cmb']:
            nent  = self._get_entries_from_json(kind=kind, trig=trigger, year=year)

            log.info(f'Taking {kind} yield from JSON as: {nent:.0f}')
        elif kind in ['bpks', 'bdks']:
            scale = self._get_rare_scale(year, trigger, kind)
            scl   = scale[0]
            nent  = nsig * scl

            log.info(f'Taking {kind} yield as {nent:.0f}={scl:.3f} * {nsig:.0f}')
        else:
            log.error(f'Invalid kind: {kind}')
            raise

        return nent
    #---------------------------------------------------------------
    def _get_signal(self, preffix, nent=None):
        year, trig = self._preffix_to_info(preffix)
        log.warning(f'Using 2018 signal for {year}')
        year = '2018'
        sig, d_cns = hqm_model.get_signal_shape(preffix=preffix, year=year, trigger=trig)

        self._add_constraints(d_cns)

        nsg    = zfit.Parameter(f'nsg_{preffix}', nent, 0, 100000)
        esig   = sig.create_extended(nsg, name='Signal')

        return esig
    #---------------------------------------------------------------
    def _get_rare(self, preffix, nent=None, kind=None):
        year, trig = self._preffix_to_info(preffix)
        log.warning(f'Using 2018 rare {kind} background for {year}')
        year = '2018'

        if   kind == 'bp':
            pdf = hqm_model.get_Bu2Ksee_shape(preffix=preffix, year=year, trigger=trig)
        elif kind == 'bd':
            pdf = hqm_model.get_Bd2Ksee_shape(preffix=preffix, year=year, trigger=trig)
        else:
            log.error(f'Invalid kind: {kind}')
            raise ValueError

        nrare= self._get_entries(nent, f'{kind}ks', trigger=trig, year=year)
        nbkg = zfit.Parameter(f'nr{kind}_{preffix}', nrare, 0, 100000)
        pdf.set_yield(nbkg)

        return pdf 
    #---------------------------------------------------------------
    def _get_prec(self, preffix, nent=None):
        year, trig = self._preffix_to_info(preffix)
        log.warning(f'Using 2018 PRec background for {year}')
        year = '2018'

        pdf = hqm_model.get_part_reco(preffix=preffix, year=year, trigger=trig)
        nent= self._get_entries(nsig=None, kind='prc', trigger=trig, year=year)
        npr = zfit.Parameter(f'npr_{preffix}', nent, 0, 100000)
        pdf.set_yield(npr)

        return pdf 
    #---------------------------------------------------------------
    def _get_msid(self, preffix):
        dset, trig    = self._get_dset_trig(preffix)
        if trig == 'GTIS':
            log.warning(f'Missing mis-id component for: {preffix}')
            return

        bld= msid(name=preffix, version='v1', obs=self._obs)
        bld.load_model()
        bld.fix_mode = 'all'
        pdf=bld.build_model() 

        return pdf
    #---------------------------------------------------------------
    def _get_pdf(self, preffix='', nent=None):
        preffix   = f'{preffix}_{self._preffix}'
        esig      = self._get_signal(preffix, nent=nent)
        self._obs = esig.space
        cbkg      = self._get_combinatorial(preffix, nent)

        if   preffix.startswith('ee_'):
            erbp  = self._get_rare(preffix, nent=nent, kind='bp')
            erbd  = self._get_rare(preffix, nent=nent, kind='bd')
            eprc  = self._get_prec(preffix, nent=nent           )
            emis  = self._get_msid(preffix)

            l_pdf = [cbkg, erbp, erbd, eprc] if emis is None else [cbkg, erbp, erbd, eprc, emis]
        elif preffix.startswith('mm_'):
            l_pdf = [cbkg]
        else:
            log.error(f'Invalid preffix: {preffix}')
            raise

        if not self._skip_signal:
            l_pdf.append(esig)
        else:
            log.warning(f'Not adding signal to fit model')

        pdf = l_pdf[0] if len(l_pdf) == 1 else zfit.pdf.SumPDF(l_pdf) 

        return pdf 
    #---------------------------------------------------------------
    def get_cons(self):
        '''
        Will return constraints on model parameters 

        Returns
        -----------
        d_val, d_var: Tuple of dictionaries pairing parameter name with value (mu of Gaussian) and 
        variance respectively.
        '''
        self._initialize()

        log.debug('-' * 20)
        log.debug(f'{"Name":<40}{"Value":<15}{"Variance":<15}')
        log.debug('-' * 20)
        for name in self._d_val:
            val = self._d_val[name]
            var = self._d_var[name]

            log.debug(f'{name:<40}{val:<15.3f}{var:<15.3f}')
        log.debug('-' * 20)

        return self._d_val, self._d_var
#---------------------------------------------------------------

