from dqc_wrapper.utils.utils_grib import *
from dqc_wrapper.utils.utils_bokeh import BokehPlot
from dqc_wrapper.unit_tests.conftest import dirs
from dqc_wrapper.test02_spatial_completeness import SpatialCompletenessTest
import pandas as pd


def prepare():
    return BokehPlot('test title', 'test-valid-ranges.html', dirs.stats_csv, dirs.hist_csv,
                     'tas', '', 0, 50, {'total': {'vmin_nmax': 0, 'vmax_nmax': 50}},
                     {'varname': 'tas', 'unit': 'degC'}, 'degC', '', '')


def test_load_grib_metadata():
    metadata = load_grib_metadata()
    assert 'grib_codes' in metadata[list(metadata.keys())[0]]


def test_get_grib_version():
    test = GribFileCheck(dirs.grib_cds)
    res = test.get_grib_version()
    assert res == 2


def test_get_grib_units():
    test = GribFileCheck(dirs.grib_cds)
    res = test.get_grib_units()
    assert res[0] == '(0 - 1)'


def test_check_mandatory_keys():
    test = GribFileCheck(dirs.grib_cds)
    err_msgs = test.check_mandatory_keys({'GRIBEditionNumber': '2'})
    assert not len(err_msgs)
    err_msgs = test.check_mandatory_keys({'GRIBEditionNumber': '27'})
    assert len(err_msgs)


def test_get_grib_first_values():
    test = GribFileCheck(dirs.grib_cds)
    res = test.get_grib_first_values()
    assert len(res)


def test_get_grib_sample():
    test = GribFileCheck(dirs.grib_cds)
    res = test.get_grib_sample()
    assert res[0][0] == 0
    assert res[1][0] == 1


def test_get_grib_hadley_regions():
    prepare()
    test = GribFileCheck(dirs.grib_cds, dt=['an'])
    hbins = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25]
    df = pd.read_csv(dirs.hist_csv)
    res = test.get_grib_hadley_regions(0, hbins, df)
    assert len(res[0]) == 15


def test_get_grib_mv():
    test = SpatialCompletenessTest([dirs.nc_masks_file], 'tas', 'nc', 'reanalysis-era5-single-levels',
                                   '', 0, dirs.masks_dir, 'yes', 'yes')
    masks = test.read_masks()
    masks[319225] = np.ones(319225)
    test = GribFileCheck(dirs.grib_cds, dt=['an'])
    res = test.get_grib_mv(masks, True)
    assert res[0] == 0
    assert res[1] == 0   # no mv or uv because both mask and grib have all non nan values

