#!/usr/bin/env python3
import unittest
from anyBSM import anyBSM
import numpy as np
import numpy.testing as test
""" This is a file with tests for 1L corrections to la112 in the THDM type II, The numbers for the numerical check were obtained in the notebook test_la112.nb analytic expressions were not included because they are too large """

class numericTest(unittest.TestCase):

    def setUp(self):
        self.THDM = anyBSM(
            anyBSM.built_in_models['THDMII'], # load the model
            scheme_name = "OSFJ", # load the OS scheme
            progress = False, # dont show any progress bars
            caching = 0, # disable the cache
            quiet=True
        )
        self.THDM.clear_cache()
        # set parameters for comparison (the check will only work for those values because it is a numerical check)
        self.THDM.setparameters({'alphaH': np.arctan(2)-np.arcsin(0.9), 'TanBeta': 2, 'Mh1': 125.1, 'Mh2': 300., 'MAh2': 300., 'MHm2': 550., 'M': 200})

    """ General 1L corrections to the 112 coupling in the THDM  (only light fermions are excluded) """
    def test_genuine_la112(self):
        """ define the particles and their moemnta """
        p = self.THDM.all_particles
        hl = p['h1']
        hh = p['h2']
        p1 = 125.1**2
        p2 = 125.1**2
        p3 = 300.**2
        """ THREE POINT B """
        hhHTPBanyH3 = self.THDM.process(hl,hl,hh,momenta=[p1,p2,p3],exclude_particles=[p for p in ['u2','u1','d3','d2','d1','u2bar','u1bar','d3bar','d2bar','d1bar','e1','e2','e3','e1bar','e2bar','e3bar']], only_topologies=['ThreePointB'])['ThreePointB']
        hhHTPBmath = -257.68536481277926+6.2071520452980025j
        test.assert_allclose(hhHTPBanyH3, hhHTPBmath)
        """ THREE POINT C """
        hhHTPCanyH3 = self.THDM.process(hl,hl,hh,momenta=[p1,p2,p3],exclude_particles=[p for p in ['u2','u1','d3','d2','d1','u2bar','u1bar','d3bar','d2bar','d1bar','e1','e2','e3','e1bar','e2bar','e3bar']], only_topologies=['ThreePointC'])['ThreePointC']
        hhHTPCmath = -75.00919028357599-8.506510393334167j
        test.assert_allclose(hhHTPCanyH3, hhHTPCmath)
        """ GENUINE (sum of the previous two) """
        hhHgenanyH3 = hhHTPBanyH3 + hhHTPCanyH3
        hhHgenmath = -332.69455509635526-2.299358348036165j
        test.assert_allclose(hhHgenanyH3, hhHgenmath)

    def test_ta_la112(self):
        """ define the particles and their momenta """
        p = self.THDM.all_particles
        hl = p['h1']
        hh = p['h2']
        p1 = 125.1**2
        p2 = 125.1**2
        p3 = 300.**2
        """ THREE POINT TA """
        hhHTPTAanyH3 = self.THDM.process(hl,hl,hh,momenta=[p1,p2,p3],exclude_particles=[p for p in ['u2','u1','d3','d2','d1','u2bar','u1bar','d3bar','d2bar','d1bar','e1','e2','e3','e1bar','e2bar','e3bar']], only_topologies=['ThreePointTA'])['ThreePointTA']
        hhHTPTAmath = 28.511383412114707
        test.assert_allclose(hhHTPTAanyH3, hhHTPTAmath)

    def test_wfra_la112(self):
        """ define the particles and their momenta """
        p = self.THDM.all_particles
        hl = p['h1']
        hh = p['h2']
        p1 = 125.1**2
        p2 = 125.1**2
        p3 = 300.**2
        """ Three point WFRA """
        hhHTPWFRanyH3 = self.THDM.process(hl,hl,hh,momenta=[p1,p2,p3],exclude_particles=[p for p in ['u2','u1','d3','d2','d1','u2bar','u1bar','d3bar','d2bar','d1bar','e1','e2','e3','e1bar','e2bar','e3bar']],only_topologies=['ThreePointWFRA'])['ThreePointWFRA'] + self.THDM.process(hl,hl,hh,momenta=[p1,p2,p3],exclude_particles=[p for p in ['u2','u1','d3','d2','d1','u2bar','u1bar','d3bar','d2bar','d1bar','e1','e2','e3','e1bar','e2bar','e3bar']],only_topologies=['ThreePointWFRB'])['ThreePointWFRB'] + self.THDM.process(hl,hl,hh,momenta=[p1,p2,p3],exclude_particles=[p for p in ['u2','u1','d3','d2','d1','u2bar','u1bar','d3bar','d2bar','d1bar','e1','e2','e3','e1bar','e2bar','e3bar']],only_topologies=['ThreePointWFRT'])['ThreePointWFRT']
        hhHTPWFRmath = 38.706861622118595-0.07839740766875303j
        test.assert_allclose(hhHTPWFRanyH3, hhHTPWFRmath)

if __name__ == '__main__':
    unittest.main()
