from dqc_wrapper.utils.utils_xarray import *
from dqc_wrapper.unit_tests.conftest import dirs
import pytest


def test_load_datafile():
    ds, _ = load_datafile(dirs.nc_file)
    assert ds['tas'].data.shape == (5, 320, 160)
    ds.close()


def test_get_variable():
    ds, _ = load_datafile(dirs.nc_file)
    response = get_variable(ds, "tas")
    assert len(response) == 5 * 320 * 160
    with pytest.raises(Exception):
        get_variable(ds, "fake_var")
    ds.close()


def test_get_mask():
    ds, _ = load_datafile(dirs.nc_file)
    response = get_mask(ds, "tas")
    assert response.shape == (320, 160)
    with pytest.raises(Exception):
        get_mask(ds, "time")
    with pytest.raises(Exception):
        get_mask(ds, "fake_var")
    ds.close()


def test_get_lat_name():
    ds, _ = load_datafile(dirs.nc_file)
    lat = get_lat_name(ds)
    assert lat == "lat"
    ds.close()


def test_get_dimensions():
    ds, _ = load_datafile(dirs.nc_file)
    assert get_dimensions(ds) == ["time"]
    ds.close()


def test_compare_matrix():
    ds, _ = load_datafile(dirs.nc_file)
    ds['tas'][0, 0, 0] = np.nan  # nan first point, values otherwise
    data = ds['tas'][0, :, :]
    mask = np.full((320, 160), 0)
    mask[-1, -1] = 1  # True last point
    mv, uv = compare_matrix(mask, data, True)
    assert mv == 0 and uv > 0
    ds.close()


def test_count_mv_uv():
    ds, _ = load_datafile(dirs.nc_file)
    ds['tas'][0, 0, 0] = np.nan  # nan first point, values otherwise
    data = ds['tas'][0, :, :]
    mask = np.full((320, 160), 0)
    mask[-1, -1] = 1  # True last point
    mv, uv = count_mv_uv(mask, data, True)
    assert mv == 0 and uv > 0
    ds.close()


def test_find_missing_unnecessary_values():
    ds, _ = load_datafile(dirs.nc_file)
    ds['tas'][0, 0, 0] = np.nan  # nan first point, values otherwise
    mask2d = ds.isel(time=0)['tas'].copy()
    mask2d[-1, -1] = 1  # False last point, values otherwise
    mv, uv = find_missing_unnecessary_values(ds, mask2d, 'tas', True)
    assert mv == 0 and uv == 0
