#! /usr/bin/env python
"""Unit tests for landlab.io.netcdf module."""
import pytest
import numpy as np
from numpy.testing import assert_array_equal

from landlab import RasterModelGrid
from landlab.io.netcdf import write_netcdf, NotRasterGridError, WITH_NETCDF4
from landlab.io.netcdf.read import _get_raster_spacing

try:
    import netCDF4 as nc
except ImportError:
    pass


def test_netcdf_write_int64_field_netcdf4(tmpdir):
    """Test write_netcdf with a grid that has an int64 field."""
    field = RasterModelGrid(4, 3)
    field.add_field('node', 'topographic__elevation',
                    np.arange(12, dtype=np.int64))

    with tmpdir.as_cwd():
        write_netcdf('test.nc', field, format='NETCDF4')

        root = nc.Dataset('test.nc', 'r', format='NETCDF4')

        for name in ['topographic__elevation']:
            assert name in root.variables
            assert_array_equal(root.variables[name][:].flatten(),
                               field.at_node[name])
            assert root.variables[name][:].dtype == 'int64'

        root.close()


def test_netcdf_write_uint8_field_netcdf4(tmpdir):
    """Test write_netcdf with a grid that has an uint8 field."""
    field = RasterModelGrid(4, 3)
    field.add_field('node', 'topographic__elevation',
                    np.arange(12, dtype=np.uint8))

    with tmpdir.as_cwd():
        write_netcdf('test.nc', field, format='NETCDF4')

        root = nc.Dataset('test.nc', 'r', format='NETCDF4')

        for name in ['topographic__elevation']:
            assert name in root.variables
            assert_array_equal(root.variables[name][:].flatten(),
                               field.at_node[name])
            assert root.variables[name][:].dtype == 'uint8'

        root.close()


def test_netcdf_write_as_netcdf3_64bit(tmpdir):
    """Test write_netcdf with output format 64-bit netcdf3."""
    from scipy.io import netcdf

    field = RasterModelGrid(4, 3)
    field.add_field('node', 'topographic__elevation', np.arange(12.))
    field.add_field('node', 'uplift_rate', 2. * np.arange(12.))

    with tmpdir.as_cwd():
        write_netcdf('test.nc', field, format='NETCDF3_64BIT')

        f = netcdf.netcdf_file('test.nc', 'r')

        for name in ['topographic__elevation', 'uplift_rate']:
            assert name in f.variables
            assert_array_equal(f.variables[name][:].flatten(), field.at_node[name])

        f.close()


def test_netcdf_write_as_netcdf3_classic(tmpdir):
    """Test write_netcdf with output format classic netcdf3."""
    from scipy.io import netcdf

    field = RasterModelGrid(4, 3)
    field.add_field('node', 'topographic__elevation', np.arange(12.))
    field.add_field('node', 'uplift_rate', 2. * np.arange(12.))

    with tmpdir.as_cwd():
        write_netcdf('test.nc', field, format='NETCDF3_CLASSIC')

        f = netcdf.netcdf_file('test.nc', 'r')

        for name in ['topographic__elevation', 'uplift_rate']:
            assert name in f.variables
            assert_array_equal(f.variables[name][:].flatten(), field.at_node[name])

        f.close()


@pytest.mark.skipif(not WITH_NETCDF4, reason='netCDF4 package not installed')
def test_netcdf_write(tmpdir):
    """Test generic write_netcdf."""
    field = RasterModelGrid(4, 3)
    field.add_field('node', 'topographic__elevation', np.arange(12.))

    with tmpdir.as_cwd():
        write_netcdf('test.nc', field, format='NETCDF4')
        root = nc.Dataset('test.nc', 'r', format='NETCDF4')

        assert set(root.dimensions) == set(['ni', 'nj', 'nt'])
        assert len(root.dimensions['ni']) == 3
        assert len(root.dimensions['nj']) == 4
        assert len(root.dimensions['nt']) == 1
        assert root.dimensions['nt'].isunlimited()

        assert set(root.variables) == set(['x', 'y', 'topographic__elevation'])

        assert_array_equal(root.variables['x'][:].flatten(),
                           np.array([0., 1., 2., 0., 1., 2., 0., 1., 2.,
                                     0., 1., 2., ]))
        assert_array_equal(root.variables['y'][:].flatten(),
                           np.array([0., 0., 0., 1., 1., 1., 2., 2., 2.,
                                     3., 3., 3., ]))
        assert_array_equal(root.variables['topographic__elevation'][:].flatten(),
                           field.at_node['topographic__elevation'])

        root.close()


@pytest.mark.skipif(not WITH_NETCDF4, reason='netCDF4 package not installed')
def test_netcdf_write_as_netcdf4_classic(tmpdir):
    """Test write_netcdf to netcdf4 classic format."""
    field = RasterModelGrid(4, 3)
    field.add_field('node', 'topographic__elevation', np.arange(12.))
    field.add_field('node', 'uplift_rate', np.arange(12.))

    with tmpdir.as_cwd():
        write_netcdf('test.nc', field, format='NETCDF4_CLASSIC')
        root = nc.Dataset('test.nc', 'r', format='NETCDF4_CLASSIC')

        for name in ['topographic__elevation', 'uplift_rate']:
            assert name in root.variables
            assert_array_equal(root.variables[name][:].flatten(),
                               field.at_node[name])

        root.close()


@pytest.mark.skipif(not WITH_NETCDF4, reason='netCDF4 package not installed')
def test_netcdf_write_names_keyword_as_list(tmpdir):
    """Test write_netcdf using a list for the *names* keyword."""
    field = RasterModelGrid(4, 3)
    field.add_field('node', 'topographic__elevation', np.arange(12.))
    field.add_field('node', 'uplift_rate', np.arange(12.))

    with tmpdir.as_cwd():
        write_netcdf('test.nc', field, names=['topographic__elevation'],
                     format='NETCDF4')
        root = nc.Dataset('test.nc', 'r', format='NETCDF4')

        assert 'topographic__elevation' in root.variables
        assert 'uplift_rate' not in root.variables
        assert_array_equal(root.variables['topographic__elevation'][:].flatten(),
                           field.at_node['topographic__elevation'])

        root.close()


@pytest.mark.skipif(not WITH_NETCDF4, reason='netCDF4 package not installed')
def test_netcdf_write_names_keyword_as_str(tmpdir):
    """Test write_netcdf using a ``str`` for the *names* keyword."""
    field = RasterModelGrid(4, 3)
    field.add_field('node', 'topographic__elevation', np.arange(12.))
    field.add_field('node', 'uplift_rate', np.arange(12.))

    with tmpdir.as_cwd():
        write_netcdf('test.nc', field, names='uplift_rate', format='NETCDF4')
        root = nc.Dataset('test.nc', 'r', format='NETCDF4')

        assert 'topographic__elevation' not in root.variables
        assert 'uplift_rate' in root.variables
        assert_array_equal(root.variables['uplift_rate'][:].flatten(),
                           field.at_node['uplift_rate'])

        root.close()


@pytest.mark.skipif(not WITH_NETCDF4, reason='netCDF4 package not installed')
def test_netcdf_write_names_keyword_as_none(tmpdir):
    """Test write_netcdf using ``None`` for the *names* keyword."""
    field = RasterModelGrid(4, 3)
    field.add_field('node', 'topographic__elevation', np.arange(12.))
    field.add_field('node', 'uplift_rate', np.arange(12.))

    with tmpdir.as_cwd():
        write_netcdf('test.nc', field, names=None, format='NETCDF4')
        root = nc.Dataset('test.nc', 'r', format='NETCDF4')

        for name in ['topographic__elevation', 'uplift_rate']:
            assert name in root.variables
            assert_array_equal(root.variables[name][:].flatten(),
                               field.at_node[name])

        root.close()


def test_2d_unit_spacing():
    """Test write_netcdf with a 2D grid with unit spacing."""
    (x, y) = np.meshgrid(np.arange(5.), np.arange(4.))

    spacing = _get_raster_spacing((y, x))
    assert spacing == 1.


def test_2d_non_unit_spacing():
    """Test _get_raster_spacing with a 2D grid with non-unit spacing."""
    (x, y) = np.meshgrid(np.arange(5.) * 2, np.arange(4.) * 2)

    spacing = _get_raster_spacing((y, x))
    assert spacing == 2.


def test_2d_uneven_spacing_axis_0():
    """Test _get_raster_spacing with a 2D grid with uneven spacing in y."""
    (x, y) = np.meshgrid(np.logspace(0., 2., num=5), np.arange(4.))

    with pytest.raises(NotRasterGridError):
        _get_raster_spacing((y, x))


def test_2d_uneven_spacing_axis_1():
    """Test _get_raster_spacing with a 2D grid with uneven spacing in x."""
    (x, y) = np.meshgrid(np.arange(4.), np.logspace(0., 2., num=5))

    with pytest.raises(NotRasterGridError):
        _get_raster_spacing((y, x))


def test_2d_switched_coords():
    """Test _get_raster_spacing with a 2D grid when the spacing is switched."""
    (x, y) = np.meshgrid(np.arange(5.), np.arange(4.))

    spacing = _get_raster_spacing((x, y))
    assert spacing == 0.


def test_1d_unit_spacing():
    """Test _get_raster_spacing with a 1D grid with unit spacing."""
    spacing = _get_raster_spacing((np.arange(5.), ))
    assert spacing == 1.


def test_1d_non_unit_spacing():
    """Test _get_raster_spacing with a 1D grid with non-unit spacing."""
    spacing = _get_raster_spacing((np.arange(5.) * 2, ))
    assert spacing == 2.


def test_1d_uneven_spacing():
    """Test _get_raster_spacing with a 1D grid with uneven spacing in y."""
    with pytest.raises(NotRasterGridError):
        _get_raster_spacing((np.logspace(0., 2., num=5), ))


@pytest.mark.skipif(not WITH_NETCDF4, reason='netCDF4 package not installed')
def test_netcdf_write_at_cells(tmpdir):
    """Test write_netcdf using with cell fields"""
    field = RasterModelGrid((4, 3))
    field.add_field('cell', 'topographic__elevation',
                    np.arange(field.number_of_cells))
    field.add_field('cell', 'uplift_rate', np.arange(field.number_of_cells))

    with tmpdir.as_cwd():
        write_netcdf('test-cells.nc', field, format='NETCDF4')
        root = nc.Dataset('test-cells.nc', 'r', format='NETCDF4')

        for name in ['topographic__elevation', 'uplift_rate']:
            assert name in root.variables
            assert_array_equal(root.variables[name][:].flatten(),
                               field.at_cell[name])

        assert set(root.dimensions) == set(['nv', 'ni', 'nj', 'nt'])
        assert len(root.dimensions['nv']) == 4
        assert len(root.dimensions['ni']) == 1
        assert len(root.dimensions['nj']) == 2
        assert len(root.dimensions['nt']) == 1
        assert root.dimensions['nt'].isunlimited()

        assert set(root.variables) == set(['x_bnds', 'y_bnds', 'topographic__elevation', 'uplift_rate'])
        root.close()
