from anyBSM.anyH3 import anyH3

class anyEffectiveCoupling(anyH3):

    _initialized = False

    def __init__(self):

        super().__init__()
        self._initialized = False
        self.init_ref_model()

    def init_ref_model(self, mref=None, reinitialize=False):
        """
        Initializes the reference model. Takes the SM by default if mref is not provided.

        Parameters:
        - mref: Optional, reference model to initialize.

        Returns:
        None
        """

        if not self._initialized or reinitialize:
            if mref is None:
                mref = anyH3('SM', scheme_name="OS", quiet=True)

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

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

        # leading corrections in ref model
        self.effWWref, self.effZZref, self.effttref, self.effbbref, self.efftataref, self.effccref, self.effssref, \
        self.effmmref, self.effuuref, self.effddref, self.effeeref = self.leading_BSM(
            mref, 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._initialized = True

    def set_particles(self, model):

        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 set_couplings(self, model, scalar):
        p = self.set_particles(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(self, 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:
        - model
        - 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 effWW, effZZ couplings
        """
        p = self.set_particles(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

    def eff_couplings(self, Qren=None, reference_model=None, printHT=False, loopWFR = False):
        """
        Computes effective couplings of neutral CP-even Higgs bosons relative to reference model
        Parameters:
        - reference_model: Optional, reference model to initialize.
        - printHT: Optional, outputs in a HiggsTools friendly format

        Returns:
        dict:  effective_couplings
        """
        if loopWFR:
            # leading corrections
            print('Warning: you are including WFR corrections to 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._initialized:
            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 = self.set_couplings(self, scalar)

                if loopWFR:
                    # leading corrections
                    effWW, effZZ, efftt, effbb, efftata, effcc, effss, effmm, effuu, effdd, effee = self.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
