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


def test_open_dataset():
   ds = open_dataset(dirs.nc_file, "")
   assert ds['tas'].data.shape == (5, 320, 160)

   ds = open_dataset(dirs.grib_file, [])
   assert ds.time.values


def test_get_variable():
    ds = open_dataset(dirs.nc_file, "")
    response = get_variable(ds, "tas")
    assert response.shape == (5, 320, 160)
    with pytest.raises(Exception):
        get_variable(ds, "fake_var")


def test_get_mask():
    ds = open_dataset(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")

def test_get_lat_name():
    ds = open_dataset(dirs.nc_file, "")
    lat = get_lat_name(ds)
    assert lat == "lat"

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


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


def test_find_missing_unnecessary_values():
    ds = open_dataset(dirs.nc_file, "").load()
    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, None, mask2d, 'tas')
    assert mv == 0 and uv == 5
