import numpy
import zfit
import math
import re
import os
import pprint
import matplotlib.pyplot as plt

from logzero     import logger  as log
from zutils.plot import plot    as zfp
from scipy.stats import poisson

#----------------------------------------------------
class model:
    def __init__(self, rk=1, preffix='', d_eff=None, d_nent=None, l_dset=None):
        '''
        Parameters
        -----------------------
        rk     (float): Value used in model, will impact toy data produced 
        preffix(str)  : Used to name parameters in case multiple models used
        d_eff  (dict): Stores full efficiency (selection + reco + acceptance) for each dataset processed
        d_nent (dict): Stores number of B -> llK entries, before selection, e.g. {r1 : 23235}
        l_dset (list): List of datasets for which the models should be created, if None, do it for all datasets i.e. r1_TOS, r2p1_TOS...
        '''
        self._obs    = zfit.Space('x', limits=(4800, 6000))
        self._rk     = rk
        self._preffix= preffix
        self._d_eff  = d_eff 
        self._d_nent = d_nent
        self._l_sub_dset = l_dset

        zfit.settings.changed_warnings.hesse_name = False

        self._l_all_dset  = ['r1', 'r2p1', '2017', '2018'] 
        self._skip_signal = None
        self._l_dset      = None
        self._d_mod       = None
        self._out_dir     = None
        self._initialized = False
    #----------------------------------------------------
    def _initialize(self):
        if self._initialized:
            return

        if self._d_eff is None:
            log.error(f'No efficiencies found, cannot provide data')
            raise

        if self._d_nent is None:
            log.error(f'No B-> Jpsi K yields found')
            raise

        self._l_dset = self._d_eff.keys() if self._l_sub_dset is None else self._l_sub_dset
        self._cache_model()

        self._initialized = True
    #----------------------------------------------------
    @property
    def out_dir(self):
        return self._out_dir

    @out_dir.setter
    def out_dir(self, value):
        try:
            os.makedirs(value, exist_ok=True)
        except:
            log.error(f'Cannot create: {value}')
            raise

        self._out_dir = value
        log.debug(f'Using output directory: {self._out_dir}')
    #----------------------------------------------------
    def _get_ds_model(self, ds, nent_mm, nent_ee):
        self._pdf_mm = self._get_pdf(preffix=f'mm_{ds}', nent=nent_mm)
        self._pdf_ee = self._get_pdf(preffix=f'ee_{ds}', nent=nent_ee)
    
        return self._pdf_mm, self._pdf_ee
    #----------------------------------------------------
    def _plot_model(self, key, mod):
        if self._out_dir is None:
            return

        plt_dir = f'{self._out_dir}/plots/models'
        os.makedirs(plt_dir, exist_ok=True)

        obj= zfp(data=mod.create_sampler(fixed_params=False), model=mod)
        obj.plot(nbins=50, ext_text=key, stacked=True)

        log.info(f'Saving to: {plt_dir}/{key}.png')
        plt.savefig(f'{plt_dir}/{key}.png')
        plt.close('all')
    #----------------------------------------------------
    def _plot_data(self, key, dat):
        if self._out_dir is None:
            return

        plt_dir = f'{self._out_dir}/plots/data'
        os.makedirs(plt_dir, exist_ok=True)

        arr_dat = dat.value().numpy()

        plt.hist(arr_dat, bins=50)

        log.info(f'Saving to: {plt_dir}/{key}.png')
        plt.title(f'{key}; {dat.name}')
        plt.savefig(f'{plt_dir}/{key}.png')
        plt.close('all')
    #----------------------------------------------------
    def _get_ds_sel_nent(self, dset):
        [ds_only, trig]= dset.split('_')
        nent           = self._d_nent[ds_only]

        eff_mm, eff_ee = self._d_eff[dset]
        nent_mm        = self._rk * nent * eff_mm
        nent_ee        =            nent * eff_ee 

        log.debug(f'{dset}:')
        log.debug(f'{nent_mm:15.0f}={self._rk:5.3e} * {nent:5.0f} * {eff_mm:5.3e}')
        log.debug(f'{nent_ee:15.0f}={       1:5.3e} * {nent:5.0f} * {eff_ee:5.3e}')
        log.debug('')

        return nent_mm, nent_ee
    #----------------------------------------------------
    def _get_selected_entries(self):
        d_nent_sel = {}

        for ds in self._l_dset:
            nent_mm = 0
            nent_ee = 0
            if ds in ['all_TOS', 'all_TIS']:
                for single_ds in [ ds.replace('all', dset) for dset in self._l_all_dset ]:
                    nmm, nee = self._get_ds_sel_nent(single_ds)
                    nent_mm += nmm
                    nent_ee += nee
            else:
                nent_mm, nent_ee = self._get_ds_sel_nent(ds)

            d_nent_sel[ds] = nent_mm, nent_ee

        return d_nent_sel
    #----------------------------------------------------
    def _cache_model(self):
        d_nent_sel = self._get_selected_entries()
        d_mod      = {}

        mod_mm_tos = None
        for ds, (nent_mm, nent_ee) in d_nent_sel.items():
            if self._l_sub_dset is not None and ds not in self._l_sub_dset:
                continue

            mod_mm, mod_ee = self._get_ds_model(ds, nent_mm, nent_ee)
            if 'TIS' in ds:
                mod_mm = mod_mm_tos
            else:
                mod_mm_tos = mod_mm

            d_mod[ds]      = mod_mm, mod_ee

            self._plot_model(f'{ds}_mm', mod_mm)
            self._plot_model(f'{ds}_ee', mod_ee)
    
        self._d_mod = d_mod
    #----------------------------------------------------
    def get_model(self, skip_signal=False):
        '''
        Returns
        -----------------
        d_ent (dict): Returns a dictionary: {name : tup} where
            name (str)  : model identifier, e.g. r1_ETOS
            tup  (tuple): Tuple with muon and electron PDFs, e.g. pdf_mm, pdf_ee 

        For muon, TIS model is the TOS one.
        '''
        self._skip_signal = skip_signal

        self._initialize()

        return self._d_mod
    #----------------------------------------------------
    def _add_ck(self, d_par):
        regex='nsg_ee_([0-9a-z]+_(TIS|TOS)_.*)'
        d_par_ck = {}
        for var_name in d_par:
            mtch = re.match(regex, var_name)
            if not mtch:
                continue

            suffix = mtch.group(1)

            ee_yld_name = var_name
            mm_yld_name = var_name.replace('_ee_', '_mm_').replace('_TIS_', '_TOS_')

            ee_yld, _   = d_par[ee_yld_name]
            mm_yld, _   = d_par[mm_yld_name]

            d_par_ck[f'ck_{suffix}'] = [ (self._rk * ee_yld) / mm_yld, 0]

        d_par.update(d_par_ck)

        return d_par
    #----------------------------------------------------
    def _add_ext_constraints(self, d_par, d_var):
        if d_var is None:
            log.warning(f'Not adding errors for constrained parameters in prefit dictionary')
            return d_par

        d_par_new = {}
        for name, var in d_var.items():
            if name not in d_par:
                log.error(f'Cannot find {name} in prefit dictionary:')
                pprint.pprint(d_par.keys())
                raise

            val = d_par[name][0]
            err = math.sqrt(var)

            d_par_new[name] = [val, err]

        d_par.update(d_par_new)

        return d_par
    #----------------------------------------------------
    def _add_ck_constraints(self, d_par, ck_cov):
        if ck_cov is None:
            log.warning(f'Not adding errors for ck parameters in prefit dictionary')
            return d_par

        d_par_new={}
        counter=0
        for name, [val, _] in d_par.items():
            if not name.startswith('ck_'):
                continue

            var = ck_cov[counter][counter]
            err = math.sqrt(var)

            d_par_new[name] = [val, err]
            counter+=1

        d_par.update(d_par_new)

        return d_par
    #----------------------------------------------------
    def get_prefit_pars(self, d_var=None, ck_cov=None):
        '''
        Used to get model parameters used to make the toy data

        Parameters
        --------------------
        d_var (dict): Dictionary with variances for parameters that are constrained. If pased the
        constraint widths will be added as errors in the prefit dictionary

        ck_cov (numpy.array): nxn numpy array representing covariance matrix for ck parameters

        Returns 
        --------------------
        d_par (dict): Dictionary storing the prefit parameters (used to build the model) and their
        errors, e.g. {'par_x' : (3, 1)}
        '''
        self._initialize()

        d_model = self.get_model()

        d_par = {}
        for mod_mm, mod_ee in d_model.values():
            d_par_mm = { par.name : [ par.value().numpy(), 0] for par in mod_mm.get_params() }
            d_par_ee = { par.name : [ par.value().numpy(), 0] for par in mod_ee.get_params() }

            d_par.update(d_par_ee)
            d_par.update(d_par_mm)

        d_par['rk'] = [self._rk, 0]
        d_par       = self._add_ck(d_par)
        d_par       = self._add_ext_constraints(d_par, d_var)
        d_par       = self._add_ck_constraints(d_par, ck_cov)

        return d_par
    #----------------------------------------------------
    def get_data(self, rseed=3):
        '''
        Returns toy data from model

        Parameters:
        -----------------
        rseed  (int):  Random seed

        Returns:
        -----------------
        d_data (dict): Dictionary with dataset and tuple of zfit data objects paired, i.e. {r1_TOS : (pdf_mm, pdf_ee) }

        For muon, TIS dataset is the TOS one.
        '''
        self._initialize()

        zfit.settings.set_seed(rseed)

        d_data     = {}
        dst_mm_tos = None
        for ds, (pdf_mm, pdf_ee) in self._d_mod.items():
            dst_mm         = pdf_mm.create_sampler(fixed_params=False)
            dst_ee         = pdf_ee.create_sampler(fixed_params=False)

            if 'TIS' in ds:
                dst_mm = dst_mm_tos
            else:
                dst_mm_tos = dst_mm

            d_data[ds]     = dst_mm, dst_ee

            self._plot_data(f'{ds}_mm', dst_mm)
            self._plot_data(f'{ds}_ee', dst_ee)

            log.debug(f'Electron data: {dst_ee.numpy().shape[0]}')
            log.debug(f'Muon data: {dst_mm.numpy().shape[0]}')

        return d_data
    #----------------------------------------------------
    @staticmethod
    def show(d_mod):
        s_dset = { key.split('_')[0] for key in d_mod }
        for dset in s_dset:
            pdf_mm_tos, pdf_ee_tos = d_mod[f'{dset}_TOS']
            pdf_mm_tis, pdf_ee_tis = d_mod[f'{dset}_TIS']

            l_par_name_mm_tos = ', '.join([ par.name for par in pdf_mm_tos.get_params() ])
            l_par_name_mm_tis = ', '.join([ par.name for par in pdf_mm_tis.get_params() ])
            l_par_name_ee_tos = ', '.join([ par.name for par in pdf_ee_tos.get_params() ])
            l_par_name_ee_tis = ', '.join([ par.name for par in pdf_ee_tis.get_params() ])

            log.info('')
            log.info(f'{dset}')
            log.info('-' * 20)
            log.info(f'{"mm TOS":<10}{l_par_name_mm_tos:<60}')
            log.info(f'{"mm TIS":<10}{l_par_name_mm_tis:<60}')
            log.info(f'{"ee TOS":<10}{l_par_name_ee_tos:<60}')
            log.info(f'{"ee TIS":<10}{l_par_name_ee_tis:<60}')
            log.info('-' * 20)
    #----------------------------------------------------
    @staticmethod
    def get_cov(kind='diag_eq', c = 0.01):
        if   kind == 'diag_eq':
            mat = numpy.diag([c] * 8)
        elif kind == 'random':
            mat = numpy.random.rand(8, 8)
            numpy.fill_diagonal(mat, 1)
            mat = mat * c
        else:
            log.error(f'Invalid kind: {kind}')
            raise
    
        return mat 
    #----------------------------------------------------
    @staticmethod
    def get_rjpsi(kind='one'):
        d_rjpsi = {}
    
        if   kind == 'one':
            d_rjpsi['d1'] = 1 
            d_rjpsi['d2'] = 1 
            d_rjpsi['d3'] = 1 
            d_rjpsi['d4'] = 1 
        elif kind == 'eff_bias':
            d_rjpsi['d1'] = 0.83333333 
            d_rjpsi['d2'] = 0.83333333 
            d_rjpsi['d3'] = 0.83333333 
            d_rjpsi['d4'] = 0.83333333 
        else:
            log.error(f'Wrong kind: {kind}')
            raise
    
        return d_rjpsi
    #----------------------------------------------------
    @staticmethod
    def get_eff(kind='equal'):
        d_eff = {}
        if   kind == 'diff':
            d_eff['d1'] = (0.6, 0.3)
            d_eff['d2'] = (0.5, 0.2)
            d_eff['d3'] = (0.7, 0.3)
            d_eff['d4'] = (0.8, 0.4)
        elif kind == 'half':
            d_eff['d1'] = (0.6, 0.3)
            d_eff['d2'] = (0.6, 0.3)
            d_eff['d3'] = (0.6, 0.3)
            d_eff['d4'] = (0.6, 0.3)
        elif kind == 'equal':
            d_eff['d1'] = (0.3, 0.3)
            d_eff['d2'] = (0.3, 0.3)
            d_eff['d3'] = (0.3, 0.3)
            d_eff['d4'] = (0.3, 0.3)
        elif kind == 'bias':
            d_eff['d1'] = (0.6, 0.25)
            d_eff['d2'] = (0.6, 0.25)
            d_eff['d3'] = (0.6, 0.25)
            d_eff['d4'] = (0.6, 0.25)
        else:
            log.error(f'Invalid kind: {kind}')
            raise
    
        return d_eff
#----------------------------------------------------

