import numpy as np
import h5py
import os
import warnings

import surfinBH

# Ignore these specific warning.
warnings.filterwarnings("ignore", message="Parameter x[[0-9]+] outside"
        " training range.")
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.")


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]

        # Download fit data if needed
        if not os.path.isfile('%s/fit_%s.h5'%(surfinBH._dataPath.DataPath(),
                short_name)):

            surfinBH.DownloadData(name)


        # Load fit
        fit = surfinBH.LoadFits(name)

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

        # Compare fit and regression data for each regression test
        test_keys = regression_h5file.keys()
        for test in test_keys:
            print('... running %s'%test)
            test_h5grp = regression_h5file[test]
            x = test_h5grp['x'].value
            fit_keys = test_h5grp['y'].keys()
            for key in fit_keys:
                print('...... testing %s'%key)
                y_reg = test_h5grp['y/%s'%key].value
                y_fit = fit(key, x)
                # 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.
                np.testing.assert_allclose(y_fit, y_reg, rtol=1e-6)

        regression_h5file.close()
