"""Tests for :mod:`esmvalcore.iris_helpers`."""
import datetime

import iris
import numpy as np
import pytest
from cf_units import Unit

from esmvalcore.iris_helpers import date2num, var_name_constraint


@pytest.fixture
def cubes():
    """Test cubes."""
    cubes = iris.cube.CubeList([
        iris.cube.Cube(0.0, var_name='a', long_name='a'),
        iris.cube.Cube(0.0, var_name='a', long_name='b'),
        iris.cube.Cube(0.0, var_name='c', long_name='d'),
    ])
    return cubes


@pytest.fixture
def units():
    return Unit('days since 0001-01-01', calendar='proleptic_gregorian')


@pytest.mark.parametrize("date, dtype, expected", [
    (datetime.datetime(1, 1, 1), np.float64, 0.0),
    (datetime.datetime(1, 1, 1), int, 0.0),
    (datetime.datetime(1, 1, 2, 12), np.float64, 1.5),
])
def test_date2num_scalar(date, dtype, expected, units):
    num = date2num(date, units, dtype=dtype)
    assert num == expected
    assert num.dtype == dtype


def test_var_name_constraint(cubes):
    """Test :func:`esmvalcore.iris_helpers.var_name_constraint`."""
    out_cubes = cubes.extract(var_name_constraint('a'))
    assert out_cubes == iris.cube.CubeList([
        iris.cube.Cube(0.0, var_name='a', long_name='a'),
        iris.cube.Cube(0.0, var_name='a', long_name='b'),
    ])
    out_cubes = cubes.extract(var_name_constraint('b'))
    assert out_cubes == iris.cube.CubeList([])
    out_cubes = cubes.extract(var_name_constraint('c'))
    assert out_cubes == iris.cube.CubeList([
        iris.cube.Cube(0.0, var_name='c', long_name='d'),
    ])
    with pytest.raises(iris.exceptions.ConstraintMismatchError):
        cubes.extract_cube(var_name_constraint('a'))
    with pytest.raises(iris.exceptions.ConstraintMismatchError):
        cubes.extract_cube(var_name_constraint('b'))
    out_cube = cubes.extract_cube(var_name_constraint('c'))
    assert out_cube == iris.cube.Cube(0.0, var_name='c', long_name='d')

[((529, 590), 'cf_units.Unit', 'Unit', (['"""days since 0001-01-01"""'], {'calendar': '"""proleptic_gregorian"""'}), "('days since 0001-01-01', calendar='proleptic_gregorian')\n", (533, 590), False, 'from cf_units import Unit\n'), ((864, 898), 'esmvalcore.iris_helpers.date2num', 'date2num', (['date', 'units'], {'dtype': 'dtype'}), '(date, units, dtype=dtype)\n', (872, 898), False, 'from esmvalcore.iris_helpers import date2num, var_name_constraint\n'), ((1093, 1117), 'esmvalcore.iris_helpers.var_name_constraint', 'var_name_constraint', (['"""a"""'], {}), "('a')\n", (1112, 1117), False, 'from esmvalcore.iris_helpers import date2num, var_name_constraint\n'), ((1317, 1341), 'esmvalcore.iris_helpers.var_name_constraint', 'var_name_constraint', (['"""b"""'], {}), "('b')\n", (1336, 1341), False, 'from esmvalcore.iris_helpers import date2num, var_name_constraint\n'), ((1367, 1389), 'iris.cube.CubeList', 'iris.cube.CubeList', (['[]'], {}), '([])\n', (1385, 1389), False, 'import iris\n'), ((1420, 1444), 'esmvalcore.iris_helpers.var_name_constraint', 'var_name_constraint', (['"""c"""'], {}), "('c')\n", (1439, 1444), False, 'from esmvalcore.iris_helpers import date2num, var_name_constraint\n'), ((1565, 1619), 'pytest.raises', 'pytest.raises', (['iris.exceptions.ConstraintMismatchError'], {}), '(iris.exceptions.ConstraintMismatchError)\n', (1578, 1619), False, 'import pytest\n'), ((1683, 1737), 'pytest.raises', 'pytest.raises', (['iris.exceptions.ConstraintMismatchError'], {}), '(iris.exceptions.ConstraintMismatchError)\n', (1696, 1737), False, 'import pytest\n'), ((1826, 1850), 'esmvalcore.iris_helpers.var_name_constraint', 'var_name_constraint', (['"""c"""'], {}), "('c')\n", (1845, 1850), False, 'from esmvalcore.iris_helpers import date2num, var_name_constraint\n'), ((1875, 1923), 'iris.cube.Cube', 'iris.cube.Cube', (['(0.0)'], {'var_name': '"""c"""', 'long_name': '"""d"""'}), "(0.0, var_name='c', long_name='d')\n", (1889, 1923), False, 'import iris\n'), ((297, 345), 'iris.cube.Cube', 'iris.cube.Cube', (['(0.0)'], {'var_name': '"""a"""', 'long_name': '"""a"""'}), "(0.0, var_name='a', long_name='a')\n", (311, 345), False, 'import iris\n'), ((355, 403), 'iris.cube.Cube', 'iris.cube.Cube', (['(0.0)'], {'var_name': '"""a"""', 'long_name': '"""b"""'}), "(0.0, var_name='a', long_name='b')\n", (369, 403), False, 'import iris\n'), ((413, 461), 'iris.cube.Cube', 'iris.cube.Cube', (['(0.0)'], {'var_name': '"""c"""', 'long_name': '"""d"""'}), "(0.0, var_name='c', long_name='d')\n", (427, 461), False, 'import iris\n'), ((650, 676), 'datetime.datetime', 'datetime.datetime', (['(1)', '(1)', '(1)'], {}), '(1, 1, 1)\n', (667, 676), False, 'import datetime\n'), ((701, 727), 'datetime.datetime', 'datetime.datetime', (['(1)', '(1)', '(1)'], {}), '(1, 1, 1)\n', (718, 727), False, 'import datetime\n'), ((745, 775), 'datetime.datetime', 'datetime.datetime', (['(1)', '(1)', '(2)', '(12)'], {}), '(1, 1, 2, 12)\n', (762, 775), False, 'import datetime\n'), ((1648, 1672), 'esmvalcore.iris_helpers.var_name_constraint', 'var_name_constraint', (['"""a"""'], {}), "('a')\n", (1667, 1672), False, 'from esmvalcore.iris_helpers import date2num, var_name_constraint\n'), ((1766, 1790), 'esmvalcore.iris_helpers.var_name_constraint', 'var_name_constraint', (['"""b"""'], {}), "('b')\n", (1785, 1790), False, 'from esmvalcore.iris_helpers import date2num, var_name_constraint\n'), ((1172, 1220), 'iris.cube.Cube', 'iris.cube.Cube', (['(0.0)'], {'var_name': '"""a"""', 'long_name': '"""a"""'}), "(0.0, var_name='a', long_name='a')\n", (1186, 1220), False, 'import iris\n'), ((1230, 1278), 'iris.cube.Cube', 'iris.cube.Cube', (['(0.0)'], {'var_name': '"""a"""', 'long_name': '"""b"""'}), "(0.0, var_name='a', long_name='b')\n", (1244, 1278), False, 'import iris\n'), ((1499, 1547), 'iris.cube.Cube', 'iris.cube.Cube', (['(0.0)'], {'var_name': '"""c"""', 'long_name': '"""d"""'}), "(0.0, var_name='c', long_name='d')\n", (1513, 1547), False, 'import iris\n')]