#! /usr/bin/env python
"""
Unit tests for landlab.io.esri_ascii module.
"""
import os

import pytest
import numpy as np
from numpy.testing import assert_array_equal, assert_array_almost_equal
from six import StringIO

from landlab.io import read_esri_ascii, read_asc_header
from landlab.io import (MissingRequiredKeyError, KeyTypeError, DataSizeError,
                        BadHeaderLineError, KeyValueError, 
                        MismatchGridDataSizeError)
from landlab import RasterModelGrid


_TEST_DATA_DIR = os.path.join(os.path.dirname(__file__), 'data')


def test_hugo_read_file_name():
    (grid, field) = read_esri_ascii(os.path.join(_TEST_DATA_DIR,
                                                 'hugo_site.asc'))

    assert isinstance(grid, RasterModelGrid)

    assert field.size == 55 * 76
    assert field.shape == (55 * 76, )


def test_hugo_read_file_like():
    with open(os.path.join(_TEST_DATA_DIR, 'hugo_site.asc')) as asc_file:
        (grid, field) = read_esri_ascii(asc_file)

    assert isinstance(grid, RasterModelGrid)

    assert field.size == 55 * 76
    assert field.shape == (55 * 76, )


def test_hugo_reshape():
    with open(os.path.join(_TEST_DATA_DIR, 'hugo_site.asc')) as asc_file:
        (grid, field) = read_esri_ascii(asc_file, reshape=True)

    assert isinstance(grid, RasterModelGrid)

    assert field.shape == (55, 76)


def test_4x3_read_file_name():
    (grid, field) = read_esri_ascii(os.path.join(_TEST_DATA_DIR,
                                                 '4_x_3.asc'))

    assert isinstance(grid, RasterModelGrid)

    assert isinstance(field, np.ndarray)
    assert_array_equal(field,
                       np.array([9., 10., 11.,
                                 6.,  7.,  8.,
                                 3.,  4.,  5.,
                                 0.,  1.,  2.]))


def test_4x3_read_file_like():
    with open(os.path.join(_TEST_DATA_DIR, '4_x_3.asc')) as asc_file:
        (grid, field) = read_esri_ascii(asc_file)

    assert isinstance(grid, RasterModelGrid)

    assert_array_equal(field,
                       np.array([9., 10., 11.,
                                 6.,  7.,  8.,
                                 3.,  4.,  5.,
                                 0.,  1.,  2.]))


def test_4x3_shape_mismatch():
    asc_file = StringIO(
        """
nrows         4
ncols         3
xllcorner     1.
yllcorner     2.
cellsize      10.
NODATA_value  -9999
1. 2. 3. 4.
5. 6. 7. 8.
9. 10. 11. 12.
        """)
    (grid, field) = read_esri_ascii(asc_file)
    assert field.size == 12

    asc_file = StringIO(
        """
nrows         4
ncols         3
xllcorner     1.
yllcorner     2.
cellsize      10.
NODATA_value  -9999
1. 2. 3. 4. 5. 6. 7. 8. 9. 10. 11. 12.
        """)
    (grid, field) = read_esri_ascii(asc_file)
    assert field.size == 12


def test_4x3_size_mismatch():
    asc_file = StringIO(
        """
nrows         4
ncols         3
xllcorner     1.
yllcorner     2.
cellsize      10.
NODATA_value  -9999
1. 2. 3. 4. 5. 6. 7. 8. 9. 10.
        """)
    with pytest.raises(DataSizeError):
        read_esri_ascii(asc_file)
    
def test_grid_data_size_mismatch():
    asc_file = StringIO(
        """
nrows         4
ncols         3
xllcorner     1.
yllcorner     2.
cellsize      10.
NODATA_value  -9999
1. 2. 3. 4. 5. 6. 7. 8. 9. 10. 11. 12.
        """)
    rmg = RasterModelGrid((10,10),10.)
    with pytest.raises(MismatchGridDataSizeError):
        read_esri_ascii(asc_file, grid=rmg)    


def test_header_missing_required_key():
    asc_file = StringIO(
        """
nrows         4
xllcorner     1.
yllcorner     2.
cellsize      10.
NODATA_value  -9999
        """)
    with pytest.raises(MissingRequiredKeyError):
        read_asc_header(asc_file)


def test_header_unknown_key():
    asc_file = StringIO(
        """
nrows         4
ncols         3
xllcorner     1.
yllcorner     2.
cellsize      10.
NODATA_value  -9999
invalid_key   1
        """)
    with pytest.raises(BadHeaderLineError):
        read_asc_header(asc_file)


def test_header_missing_value():
    asc_file = StringIO(
        """
nrows         4
ncols         3
xllcorner     1.
yllcorner     2.
cellsize
NODATA_value  -9999
invalid_key   1
        """)
    with pytest.raises(BadHeaderLineError):
        read_asc_header(asc_file)


def test_header_bad_values():
    asc_file = StringIO(
        """
nrows         -4
ncols         3
xllcorner     1.
yllcorner     2.
cellsize      10.
NODATA_value  -9999
        """)
    with pytest.raises(KeyValueError):
        read_asc_header(asc_file)


def test_header_missing_mutex_key():
    asc_file = StringIO(
        """
ncols         3
nrows         4
yllcorner     2.
cellsize      10.
NODATA_value  -9999
        """)
    with pytest.raises(MissingRequiredKeyError):
        read_asc_header(asc_file)


def test_header_mutex_key():
    asc_file = StringIO(
        """
ncols         3
nrows         4
xllcenter     1.
yllcorner     2.
cellsize      10.
NODATA_value  -9999
        """)
    header = read_asc_header(asc_file)
    assert header['xllcenter'] == 1.
    with pytest.raises(KeyError):
        header['xllcorner']

    asc_file = StringIO(
        """
ncols         3
nrows         4
xllcorner     1.
yllcorner     2.
cellsize      10.
NODATA_value  -9999
        """)
    header = read_asc_header(asc_file)
    assert header['xllcorner'] == 1.
    with pytest.raises(KeyError):
        header['xllcenter']


def test_header_missing_optional():
    asc_file = StringIO(
        """
ncols         3
nrows         4
xllcenter     1.
yllcorner     2.
cellsize      10.
        """)
    header = read_asc_header(asc_file)
    with pytest.raises(KeyError):
        header['nodata_value']


def test_header_case_insensitive():
    asc_file = StringIO(
        """
nCoLs         3
nrows         4
Xllcenter     1.
YLLCORNER     2.
CELLSIZE      10.
NODATA_value  -999
        """)
    header = read_asc_header(asc_file)
    for key in ['ncols', 'nrows', 'xllcenter', 'yllcorner', 'cellsize',
                'nodata_value']:
        assert key in header


def test_header_wrong_type():
    asc_file = StringIO(
        """
nCoLs         3.5
nrows         4
Xllcenter     1.
YLLCORNER     2.
CELLSIZE      10.
NODATA_value  -999
        """)
    with pytest.raises(KeyTypeError):
        read_asc_header(asc_file)


def test_name_keyword():
    (grid, field) = read_esri_ascii(os.path.join(_TEST_DATA_DIR,
                                                 '4_x_3.asc'),
                                    name='air__temperature')

    assert isinstance(grid, RasterModelGrid)

    assert isinstance(field, np.ndarray)
    assert_array_equal(field,
                       np.array([9., 10., 11.,
                                 6.,  7.,  8.,
                                 3.,  4.,  5.,
                                 0.,  1.,  2.]))
    assert_array_almost_equal(grid.at_node['air__temperature'], field)
    assert grid.at_node['air__temperature'] is field
    
def test_halo_keyword():
    (grid, field) = read_esri_ascii(os.path.join(_TEST_DATA_DIR, \
                                                 '4_x_3.asc'), \
                                                 halo=1)
                                    
    assert isinstance(grid, RasterModelGrid)

    assert isinstance(field, np.ndarray)
    assert_array_equal(field,
                       np.array([-9999., -9999., -9999., -9999., -9999.,  
                                 -9999.,     9.,    10.,    11., -9999.,
                                 -9999.,     6.,     7.,     8., -9999.,
                                 -9999.,     3.,     4.,     5., -9999.,
                                 -9999.,     0.,     1.,     2., -9999.,
                                 -9999., -9999., -9999., -9999., -9999.]))
                                 
def test_halo_keyword_no_nodata_value():
    (grid, field) = read_esri_ascii(os.path.join(_TEST_DATA_DIR, \
                                                '4_x_3_no_nodata_value.asc'), \
                                                halo=1)
                                    
    assert isinstance(grid, RasterModelGrid)

    assert isinstance(field, np.ndarray)
    assert_array_equal(field,
                       np.array([-9999., -9999., -9999., -9999., -9999.,  
                                 -9999.,     9.,    10.,    11., -9999.,
                                 -9999.,     6.,     7.,     8., -9999.,
                                 -9999.,     3.,     4.,     5., -9999.,
                                 -9999.,     0.,     1.,     2., -9999.,
                                 -9999., -9999., -9999., -9999., -9999.]))


if __name__ == '__main__':
    unittest.main()
