import logging
from anyBSM.anyH3 import anyH3

logging.basicConfig(format='%(levelname)s:%(message)s')
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())

__doc__ = r"""
<div class="mermaid">
flowchart TB

  subgraph classes
    anyH3 --> anyEffectiveCoupling;
  end
  subgraph methods
    anyEffectiveCoupling --> init_ref_model --> eff_couplings;
  end
</div>
"""

class anyEffectiveCoupling(anyH3):
    r""" Compute and format single-Higgs-coupling modifiers

    Example:
    ```python
    from anyBSM import anyBSM
    THDM = anyBSM('THDMII')
    THDM.eff_couplings(
        reference_model = 'SM', # (SM=default) construct coupling modifier w.r.t. this model
        printHT = True, # return a dict that can be directly passed to HiggsTools
        loopWFR = True, # one-loop take corrections to the Higgs-leg into account (MSbar)
        Qren = 173. # renormalisation scheme used in the one-loop corrections
        )
    ```
    """

    def check_init_ref_model(self):
        """ checks if a reference model was already initialized """
        if not hasattr(self, '_initialized'):
            self.ref_initialized = False

        return self.ref_initialized

    def init_ref_model(self, model = None, force :  bool = False):
        """
        Initializes the couplings of the reference model. Takes the SM by default if `model` was not provided.

        Args:
            model (`anyBSM.anyModel`): model-object to use as reference model. If not provided, the 'SM' model is initialized and used
            force (bool): force re-initialization of the effective couplings of the model

        Returns:
            None
        """
        if self.check_init_ref_model() and not force:
            return

        if model is None:
            model = anyH3('SM', scheme_name="OS", quiet=True)

        # particles in the ref. model
        self.pref = get_SMparticles(model)

        # couplings in ref model
        self.cplsref = get_singleHiggs_couplings(model, model.SM_particles['Higgs-Boson'])

        # leading corrections in ref model
        self.effWWref, self.effZZref, self.effttref, self.effbbref, self.efftataref, self.effccref, self.effssref, \
        self.effmmodel, self.effuuref, self.effddref, self.effeeref = leading_BSM(
            model, self.cplsref['cw'], self.cplsref['cz'], self.cplsref['yt1'],
            self.cplsref['yb1'], self.cplsref['ytau1'], self.cplsref['yc1'], self.cplsref['ys1'], self.cplsref['ym1'],
            self.cplsref['yu1'], self.cplsref['yd1'], self.cplsref['ye1'], Qren=None
        )

        self.ref_initialized = True

    def eff_couplings(self, reference_model=None, printHT=False, loopWFR = False, Qren=None):
        """
        Computes effective couplings of neutral CP-even Higgs bosons relative to the reference model `reference_model`.
        Args:
            reference_model (anyBSM.anyModel): reference model to use
            printHT (bool): returns dictionary in a HiggsTools friendly format (default=False)
            loopWFR (bool): whether to dress the external Higgs-leg with one-loop (MSbar) Z-factors. Additionally, the trilinear Higgs coupling is computed at the full one-loop level using `anyBSM.anyH3`.
            Qren (float): renormalisation scale used in the loopWFR-option (default: mean of all masses heavyier than 150GeV)

        Returns:
            dict:  effective couplings normalised to the corresponding effective couplings in the reference model
        """
        if loopWFR:
            # leading corrections
            logger.warning('You are including WFR corrections to the Higgs leg in all couplings')

        if reference_model is not None: # this will reinitialize the model everytime reference_model is true
            self.init_ref_model(mref=reference_model, reinitialize=True)

        if not self.check_init_ref_model():
            self.init_ref_model()

        # couplings to fermions
        scalars = [s for s in self.particles['S'] if (not (s.goldstone or s.anti().goldstone)) and s.charge == 0]

        effective_couplings = {}
        n = 25
        for scalar in scalars:
            try: # this is to avoid taking into account pseudoscalars
                cpls = get_singleHiggs_couplings(self, scalar)

                if loopWFR:
                    # leading corrections
                    effWW, effZZ, efftt, effbb, efftata, effcc, effss, effmm, effuu, effdd, effee = leading_BSM(
                        self, cpls['cw'], cpls['cz'], cpls['yt1'],
                        cpls['yb1'], cpls['ytau1'], cpls['yc1'], cpls['ys1'], cpls['ym1'],
                        cpls['yu1'], cpls['yd1'], cpls['ye1'], Qren
                    )
                    couphhh = cpls['hhh']['total'].real / self.cplsref['hhh']['treelevel'].real
                    coupZZ = (effZZ / self.cplsref['cz']).real
                    coupWW = (effWW / self.cplsref['cw']).real
                    coupdde = (effdd / self.cplsref['yd1']).real
                    coupuue = (effuu / self.cplsref['yu1']).real
                    coupsse = (effss / self.cplsref['ys1']).real
                    coupcce = (effcc / self.cplsref['yc1']).real
                    coupbbe = (effbb / self.cplsref['yb1']).real
                    couptte = (efftt / self.cplsref['yt1']).real
                    coupeee = (effee / self.cplsref['ye1']).real
                    coupmme = (effmm / self.cplsref['ym1']).real
                    couptatae = (efftata / self.cplsref['ytau1']).real
                    coupddo = cpls['yd5']
                    coupuuo = cpls['yu5']
                    coupsso = cpls['ys5']
                    coupcco = cpls['yc5']
                    coupbbo = cpls['yb5']
                    couptto = cpls['yt5']
                    coupeeo = cpls['ye5']
                    coupmmo = cpls['ym5']
                    couptatao = cpls['ytau5']

                else:
                    couphhh = cpls['hhh']['treelevel'].real / self.cplsref['hhh']['treelevel'].real
                    coupZZ = (cpls['cz'] / self.cplsref['cz']).real
                    coupWW = (cpls['cw'] / self.cplsref['cw']).real
                    coupdde = (cpls['yd1'] / self.cplsref['yd1']).real
                    coupuue = (cpls['yu1'] / self.cplsref['yu1']).real
                    coupsse = (cpls['ys1'] / self.cplsref['ys1']).real
                    coupcce = (cpls['yc1'] / self.cplsref['yc1']).real
                    coupbbe = (cpls['yb1'] / self.cplsref['yb1']).real
                    couptte = (cpls['yt1'] / self.cplsref['yt1']).real
                    coupeee = (cpls['ye1'] / self.cplsref['ye1']).real
                    coupmme = (cpls['ym1'] / self.cplsref['ym1']).real
                    couptatae = (cpls['ytau1'] / self.cplsref['ytau1']).real
                    coupddo = cpls['yd5']
                    coupuuo = cpls['yu5']
                    coupsso = cpls['ys5']
                    coupcco = cpls['yc5']
                    coupbbo = cpls['yb5']
                    couptto = cpls['yt5']
                    coupeeo = cpls['ye5']
                    coupmmo = cpls['ym5']
                    couptatao = cpls['ytau5']

                # mass and width
                mH = scalar.mass
                width = scalar.width.nvalue
                if width == 0.0:
                    width = -(self.Sigma(scalar, momentum=f'{mH}**2')/mH.nvalue).imag

                # trilinear coupling for the 125 GeV Higgs
                if n == 25:
                    if printHT:
                        effective_couplings[f'effc_{n}_lam'] = couphhh
                    else:
                        effective_couplings[f'kappa_lambda_{n}'] = couphhh

                # result in HT format
                if printHT:
                    effective_couplings.update({
                        f'm_{n}': mH.value,
                        f'w_{n}': width,
                        f'effc_{n}_gg': 1,
                        f'effc_{n}_gamgam': 1,
                        f'effc_{n}_ZZ': coupZZ,
                        f'effc_{n}_WW': coupWW,
                        f'effc_{n}_dd_s': coupdde,
                        f'effc_{n}_uu_s': coupuue,
                        f'effc_{n}_ss_s': coupsse,
                        f'effc_{n}_cc_s': coupcce,
                        f'effc_{n}_bb_s': coupbbe,
                        f'effc_{n}_tt_s': couptte,
                        f'effc_{n}_ee_s': coupeee,
                        f'effc_{n}_mumu_s': coupmme,
                        f'effc_{n}_tautau_s': couptatae,
                        f'effc_{n}_dd_p': coupddo,
                        f'effc_{n}_uu_p': coupuuo,
                        f'effc_{n}_ss_p': coupsso,
                        f'effc_{n}_cc_p': coupcco,
                        f'effc_{n}_bb_p': coupbbo,
                        f'effc_{n}_tt_p': couptto,
                        f'effc_{n}_ee_p': coupeeo,
                        f'effc_{n}_mumu_p': coupmmo,
                        f'effc_{n}_tautau_p': couptatao,
                        f'CP_{n}': 1
                    })

                # result
                else:
                    effective_couplings.update({
                        f'kappa_Z_{n}': coupZZ,
                        f'kappa_W_{n}': coupWW,
                        f'kappa_t_even_{n}': couptte,
                        f'kappa_t_odd_{n}': couptto,
                        f'kappa_b_even_{n}': coupbbe,
                        f'kappa_b_odd_{n}': coupbbo,
                        f'kappa_tau_even_{n}': couptatae,
                        f'kappa_tau_odd_{n}': couptatao,
                        f'kappa_c_even_{n}': coupcce,
                        f'kappa_c_odd_{n}': coupcco,
                        f'kappa_s_even_{n}': coupsse,
                        f'kappa_s_odd_{n}': coupsso,
                        f'kappa_mu_even_{n}': coupmme,
                        f'kappa_mu_odd_{n}': coupmmo,
                        f'kappa_u_even_{n}': coupuue,
                        f'kappa_u_odd_{n}': coupuuo,
                        f'kappa_d_even_{n}': coupdde,
                        f'kappa_d_odd_{n}': coupddo,
                        f'kappa_e_even_{n}': coupeee,
                        f'kappa_e_odd_{n}': coupeeo
                    })
                n += 10
            except:
                pass
        return effective_couplings

def get_SMparticles(model):
    """ introduce abbreviations for SM particles.
    It's a short-hand for accessing `anyBSM.anyModel.SM_particles`

    Args:
        model (`anyBSM.anyModel`): model-object with initialized SM_particles

    Returns:
        dict: with keys 'H', 'Z', 'W', 'T', 'C', 'U', 'B', 'S', 'D', 'Tau', 'M', 'E', 'A' and values pointing to the UFO particle objects.
    """

    particles = {
    'H' : model.SM_particles['Higgs-Boson'],
    'Z' : model.SM_particles['Z-Boson'],
    'W' : model.SM_particles['W-Boson'],
    'T' : model.SM_particles['Top-Quark'],
    'C' : model.SM_particles['Charm-Quark'],
    'U' : model.SM_particles['Up-Quark'],
    'B' : model.SM_particles['Bottom-Quark'],
    'S' : model.SM_particles['Strange-Quark'],
    'D' : model.SM_particles['Down-Quark'],
    'Tau' : model.SM_particles['Tau-Lepton'],
    'M' : model.SM_particles['Muon-Lepton'],
    'E' : model.SM_particles['Electron-Lepton'],
    'A' : model.SM_particles['Photon'],
    }

    return particles

def get_singleHiggs_couplings(model, scalar):
    """ introduce abbreviations for single-Higgs couplings and evaluate then numerically

    Args:
        model (`anyBSM.anyModel`): model-object to use
        scalar (`anyBSM.ufo.object_library.Particle`): scalar to use

    Returns:
        dict: with keys `'yt1', 'yt5', 'yc1', 'yc5', 'yu1', 'yu5', 'yb1', 'yb5', 'ys1', 'ys5', 'yd1', 'yd5', 'ytau1', 'ytau5', 'ym1', 'ym5', 'ye1', 'ye5', 'cz', 'cw', 'hhh'` where 1/5 are for the corresponding Lorentz structure.
    """

    p = get_SMparticles(model)
    couplings = {}

    for fermion in ['T', 'C', 'U', 'B', 'S', 'D', 'Tau', 'M', 'E']:
        couplings[f'y{fermion.lower()}1'] = model._eval(model.getcoupling15(scalar, p[fermion], p[fermion].anti())['1'])
        couplings[f'y{fermion.lower()}5'] = model._eval(model.getcoupling15(scalar, p[fermion], p[fermion].anti())['5'])

    couplings.update({
        'cz': model._eval(model.getcoupling15(scalar, p['Z'], p['Z'])['c']),
        'cw': model._eval(model.getcoupling15(scalar, p['W'], p['W'].anti())['c']),
        'hhh': model.lambdahhh()
    })
    return couplings

def leading_BSM(model, cw, cz, ct, cb, cta, cc, cs, cm, cu, cd, ce, Qren):
    """
    Computes leading scalar corrections to the couplings ghZZ and ghWW
    including only wave function renormalization in an MSbar at a Qren scale.
    Args:
        model (anyBSM.anyModel): model objects
        cw, cz: tree level ghZZ and ghWW couplings
        ct, cb, cta, cc, cs, cm, cu, cd, ce: tree level SFF couplings
        Qren: renormalization scale for MSbar scheme, if None then it takes the average of the masses of the heavy particles in the model

    Returns:
        one loop effective single-Higgs couplings
        `effWW, effZZ, efftt, effbb, efftata, effcc, effss, effmm, effuu, effdd, effee`
    """

    p = get_SMparticles(model)

    scalars = [s for s in model.particles['S'] if (not (s.goldstone or s.anti().goldstone)) and s.charge == 0]

    effWW = cw
    effZZ = cz
    efftt = ct
    effbb = cb
    efftata = cta
    effcc = cc
    effss = cs
    effmm = cm
    effuu = cu
    effdd = cd
    effee = ce

    masses = [s.mass.nvalue for s in model.all_particles.values() if s.mass.nvalue > 150]
    if Qren is None:
        Qren = sum(masses) / len(masses)

    model.setparameters({'Qren': Qren})
    for x in scalars:
        if x == p['H']:
            mH = p['H'].mass
            dZ11 = model.Sigmaprime(p['H'], momentum=f'{mH}**2')
            effWW += dZ11/2 * cw
            effZZ += dZ11/2 * cz
            efftt += dZ11/2 * ct
            effbb += dZ11/2 * cb
            efftata += dZ11/2 * cta
            effcc += dZ11/2 * cc
            effss += dZ11/2 * cs
            effmm += dZ11/2 * cm
            effuu += dZ11/2 * cu
            effdd += dZ11/2 * cd
            effee += dZ11/2 * ce
        else:
            try:
                cxz = model._eval(model.getcoupling15(x, p['Z'], p['Z'])['c'])
                cxw = model._eval(model.getcoupling15(x, p['W'], p['W'].anti())['c'])
                cxt = model._eval(model.getcoupling15(x, p['T'], p['T'].anti())['1'])
                cxb = model._eval(model.getcoupling15(x, p['B'], p['B'].anti())['1'])
                cxta = model._eval(model.getcoupling15(x, p['Tau'], p['Tau'].anti())['1'])
                cxc = model._eval(model.getcoupling15(x, p['C'], p['C'].anti())['1'])
                cxs = model._eval(model.getcoupling15(x, p['S'], p['S'].anti())['1'])
                cxm = model._eval(model.getcoupling15(x, p['M'], p['M'].anti())['1'])
                cxu = model._eval(model.getcoupling15(x, p['U'], p['U'].anti())['1'])
                cxd = model._eval(model.getcoupling15(x, p['D'], p['D'].anti())['1'])
                cxe = model._eval(model.getcoupling15(x, p['E'], p['E'].anti())['1'])

                mX = x.mass
                Sigxh = model.Sigma(x, p['H'], momentum=f'{mH}**2')
                dZ21 = Sigxh / (mH.nvalue**2 - mX.nvalue**2)

                effWW += dZ21 * cxw
                effZZ += dZ21 * cxz
                efftt += dZ21 * cxt
                effbb += dZ21 * cxb
                efftata += dZ21 * cxta
                effcc += dZ21 * cxc
                effss += dZ21 * cxs
                effmm += dZ21 * cxm
                effuu += dZ21 * cxu
                effdd += dZ21 * cxd
                effee += dZ21 * cxe
            except:
                pass

    return effWW, effZZ, efftt, effbb, efftata, effcc, effss, effmm, effuu, effdd, effee
