from extractor import extractor as ext

import os
import numpy
import zfit
import math
import pprint
import pytest
import utils_noroot as utnr
import rk.utilities as rkut

from logzero    import logger    as log
from np_reader  import np_reader as np_rdr
from rkex_model import model

#-----------------------------
def skip_test():
    try:
        uname = os.environ['USER']
    except:
        pytest.skip()

    if uname in ['angelc', 'campoverde']:
        return
#----------------------------------------------------
def check_close(value, target, abs_tol=1e-4):
    pas_val = math.isclose(value, target, abs_tol = abs_tol)
    if not pas_val:
        log.error(f'{value:.6f} != {target:.6f}')
        raise
#----------------------------------------------------
def test_simple():
    log.info('Running: test_simple')
    d_eff       = model.get_eff()
    d_yld       = {'d1' : 2e3, 'd2' : 2e3, 'd3' : 2e3, 'd4' : 2e3} 

    mod         = model(preffix='simple', d_eff=d_eff)
    d_dat       = mod.get_data(d_nent=d_yld)
    d_mod       = mod.get_model()

    obj         = ext()
    obj.eff     = d_eff 
    obj.data    = d_dat
    obj.model   = d_mod 
    obj.plt_dir = 'tests/extractor/simple'

    result      = obj.get_fit_result()
    result.hesse()

    d_rk   = result.params['rk']
    rk_val = d_rk['value']
    rk_err = d_rk['hesse']['error']
#----------------------------------------------------
def test_efficiency():
    log.info('Running: test_efficiency')
    d_eff        = model.get_eff(kind='diff')
    d_yld        = {'d1' : 2e3, 'd2' : 2e3, 'd3' : 2e3, 'd4' : 2e3} 

    mod          = model(preffix='efficiency', d_eff=d_eff)
    d_dat        = mod.get_data(d_nent=d_yld)
    d_mod        = mod.get_model()

    obj          = ext()
    obj.eff      = d_eff 
    obj.data     = d_dat
    obj.model    = d_mod 
    obj.plt_dir  = 'tests/extractor/efficiency'

    result       = obj.get_fit_result()
    result.hesse()

    d_rk   = result.params['rk']
    rk_val = d_rk['value']
    rk_err = d_rk['hesse']['error']
#----------------------------------------------------
def test_constraint():
    log.info('Running: test_constraint')
    d_eff        = model.get_eff(kind='half')
    d_yld        = {'d1' : 2e3, 'd2' : 2e3, 'd3' : 2e3, 'd4' : 2e3} 

    mod          = model(preffix='constraint', d_eff=d_eff)
    cvmat        = mod.get_cov(kind='diag_eq', c=0.001)
    d_dat        = mod.get_data(d_nent=d_yld)
    d_mod        = mod.get_model()

    obj          = ext()
    obj.eff      = d_eff 
    obj.cov      = cvmat
    obj.data     = d_dat
    obj.model    = d_mod 
    obj.plt_dir  = 'tests/extractor/constraint'

    result       = obj.get_fit_result()

    result.hesse()

    d_rk   = result.params['rk']
    rk_val = d_rk['value']
    rk_err = d_rk['hesse']['error']
#----------------------------------------------------
def test_rjpsi():
    log.info('Running: test_rjpsi')
    d_eff        = model.get_eff(kind='half')
    d_yld        = {'d1' : 2e3, 'd2' : 2e3, 'd3' : 2e3, 'd4' : 2e3} 

    mod          = model(preffix='rjpsi', d_eff=d_eff)
    cvmat        = mod.get_cov(kind='diag_eq', c=0.001)
    d_rjpsi      = mod.get_rjpsi(kind='eff_bias')
    d_dat        = mod.get_data(d_nent=d_yld)
    d_mod        = mod.get_model()

    obj          = ext()
    obj.rjpsi    = d_rjpsi
    obj.eff      = d_eff
    obj.cov      = cvmat
    obj.data     = d_dat
    obj.model    = d_mod 
    obj.plt_dir  = 'tests/extractor/rjpsi'

    result       = obj.get_fit_result()

    result.hesse()

    d_rk   = result.params['rk']
    rk_val = d_rk['value']
    rk_err = d_rk['hesse']['error']
#----------------------------------------------------
def test_real():
    skip_test()
    log.info('Running: test_real')

    rdr          = np_rdr(sys='v65', sta='v63', yld='v24')
    rdr.cache    = True
    cv_sys       = rdr.get_cov(kind='sys')
    cv_sta       = rdr.get_cov(kind='sta')
    d_eff        = rdr.get_eff()
    d_rjpsi      = rdr.get_rjpsi()
    d_byld       = rdr.get_byields()
    d_nent       = rkut.average_byields(d_byld, l_exclude=['TIS'])
    d_rare_yld   = rkut.reso_to_rare(d_nent, kind='jpsi')

    mod          = model(preffix='real', d_eff=d_eff)
    d_mod        = mod.get_model()
    d_dat        = mod.get_data(d_nent=d_rare_yld)

    obj          = ext()
    obj.rjpsi    = d_rjpsi
    obj.eff      = d_eff
    obj.cov      = cv_sys + cv_sta
    obj.data     = d_dat
    obj.model    = d_mod 
    obj.plt_dir  = 'tests/extractor/real'
    result       = obj.get_fit_result()

    log.info(f'Calculating errors')
    result.hesse()
    result.freeze()
    utnr.dump_pickle(result, 'tests/extractor/real/result.pkl')
#----------------------------------------------------
def main():
    test_real()
    test_rjpsi()
    test_constraint()
    test_simple()
    test_efficiency()
#----------------------------------------------------
if __name__ == '__main__':
    main()

