import numpy as np
import h5py
import os
import warnings

import surfinBH

# Ignore these specific warnings.
warnings.filterwarnings("ignore", message="Mass ratio outside training range.")
warnings.filterwarnings("ignore", message="Spin magnitude of BhA outside"
        " training range.")
warnings.filterwarnings("ignore", message="Spin magnitude of BhB outside"
        " training range.")
warnings.filterwarnings("ignore", message="Extrapolating dynamics to")

# lower rtol seems to make the test fail for the GPR error
# estimates when testing on a different machine. This
# seems to happen even with the same version of sklearn.
rtol = 1e-6

def test_fit_regression():
    """ Compares all existing fits against saved regression data.
        Regression data should already have been generated by
        generate_regression_data.py
    """
    # List of all available fits
    fit_names = surfinBH.fits_collection.keys()
    for name in fit_names:
        short_name = name.split('surfinBH')[-1]

        # Load fit
        fit = surfinBH.LoadFits(name)

        # Load regression data
        regression_h5file = h5py.File('test/regression_data/fit_%s.h5'%(
                short_name))

        extra_kwargs_list = fit._extra_regression_kwargs()

        # Each group in the h5file can have different kwargs to test
        kwargs_grp_keys = regression_h5file.keys()
        for kw_grp_key in kwargs_grp_keys:
            kw_h5grp = regression_h5file[kw_grp_key]

            print('\nrunning %s'%kw_grp_key)

            if kw_grp_key == 'No_kwargs':
                kwargs = {}
            else:
                kwargs = extra_kwargs_list[int(kw_grp_key.split('_set_')[-1])]

            # Compare fit and regression data for each regression test
            test_keys = kw_h5grp.keys()
            for test in test_keys:
                print('... running %s'%test)
                test_h5grp = kw_h5grp[test]
                q = test_h5grp['q'].value
                chiA = test_h5grp['chiA'].value
                chiB = test_h5grp['chiB'].value

                # remnant mass
                y_reg = test_h5grp['y/mf'].value
                y_fit = fit.mf(q, chiA, chiB, **kwargs)
                np.testing.assert_allclose(y_fit, y_reg, rtol=rtol)

                # remnant spin
                y_reg = test_h5grp['y/chif'].value
                y_fit = fit.chif(q, chiA, chiB, **kwargs)
                np.testing.assert_allclose(y_fit, y_reg, rtol=rtol)

                # remnant kick
                y_reg = test_h5grp['y/vf'].value
                y_fit = fit.vf(q, chiA, chiB, **kwargs)
                np.testing.assert_allclose(y_fit, y_reg, rtol=rtol)

        regression_h5file.close()
