#!/usr/bin/env pytest
###############################################################################
# $Id$
#
# Project:  GDAL/OGR Test Suite
# Purpose:  Test multidimensional support in HDF4 driver
# Author:   Even Rouault <even.rouault@spatialys.com>
#
###############################################################################
# Copyright (c) 2019, Even Rouault <even.rouault@spatialys.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
###############################################################################

import gdaltest
import pytest
import struct

from osgeo import gdal

pytestmark = pytest.mark.require_driver('HDF4')


###############################################################################
# Test reading HDFEOS SWATH products


def test_hdf4multidim_hdfeos_swath():

    if not gdaltest.download_file('https://gamma.hdfgroup.org/ftp/pub/outgoing/NASAHDFfiles2/eosweb/hdf4/hdfeos2-swath-wo-dimmaps/AMSR_E_L2_Ocean_B01_200206182340_A.hdf', 'AMSR_E_L2_Ocean_B01_200206182340_A.hdf'):
        pytest.skip()

    ds = gdal.OpenEx('tmp/cache/AMSR_E_L2_Ocean_B01_200206182340_A.hdf', gdal.OF_MULTIDIM_RASTER)
    assert ds
    rg = ds.GetRootGroup()
    assert rg
    assert rg.GetGroupNames() == ['swaths']
    attrs = rg.GetAttributes()
    assert attrs
    attr = rg.GetAttribute('HDFEOSVersion')
    assert attr.Read() == 'HDFEOS_V2.7.2'
    swaths = rg.OpenGroup('swaths')
    assert swaths
    assert not rg.OpenGroup('foo')
    assert swaths.GetGroupNames() == ['Swath1']
    swath1 = swaths.OpenGroup('Swath1')
    assert swath1
    assert not swaths.OpenGroup('foo')
    attrs = swath1.GetAttributes()
    assert len(attrs) == 3
    attr = swath1.GetAttribute('SoftwareRevisionDate')
    assert attr.Read()[0:-1] == "November 7, 2003"
    attr = swath1.GetAttribute('SoftwareBuildNumber')
    assert attr.Read() == 1
    dims = swath1.GetDimensions()
    assert len(dims) == 4
    dim = next((x for x in dims if x.GetName() == 'DataXtrack_lo'), None)
    assert dim
    assert dim.GetFullName() == '/swaths/Swath1/DataXtrack_lo'
    assert dim.GetSize() == 196 
    assert swath1.GetGroupNames() == ['Data Fields', 'Geolocation Fields']
    assert not swath1.OpenGroup('foo')
    datafields = swath1.OpenGroup('Data Fields')
    assert datafields
    assert len(datafields.GetMDArrayNames()) == 10
    assert not datafields.OpenMDArray('foo')
    array = datafields.OpenMDArray('High_res_cloud')
    assert array
    dims = array.GetDimensions()
    assert len(dims) == 2
    assert dims[0].GetFullName() == '/swaths/Swath1/DataTrack_lo'
    assert array.GetDataType().GetNumericDataType() == gdal.GDT_Int16
    attr = array.GetAttribute('Scale')
    assert attr.Read() == 9.999999747378752e-05

    got_data = array.Read(array_start_idx = [13, 0], count = [3, 2])
    assert len(got_data) == 3 * 2 * 2
    assert struct.unpack('h' * 6, got_data) == (0, 0,
                                                17318, 17317,
                                                17318, 17317)


    got_data = array.Read(array_start_idx = [13, 0], count = [3, 2],
                          buffer_datatype = gdal.ExtendedDataType.Create(gdal.GDT_Int32))
    assert len(got_data) == 3 * 2 * 4
    assert struct.unpack('i' * 6, got_data) == (0, 0,
                                                17318, 17317,
                                                17318, 17317)

    got_data = array.Read(array_start_idx = [13, 0], count = [3, 2], array_step = [2, 1])
    assert len(got_data) == 3 * 2 * 2
    assert struct.unpack('h' * 6, got_data) == (0, 0,
                                                17318, 17317,
                                                17317, 17317)

    got_data = array.Read(array_start_idx = [15, 1], count = [3, 2], array_step = [-1, -1])
    assert len(got_data) == 3 * 2 * 2
    assert struct.unpack('h' * 6, got_data) == (17317, 17318,
                                                17317, 17318,
                                                0, 0)

    got_data = array.Read(array_start_idx = [13, 0], count = [3, 2], buffer_stride = [1, 3])
    assert len(got_data) == 3 * 2 * 2
    assert struct.unpack('h' * 6, got_data) == (0, 17318, 17318,
                                                0, 17317, 17317)

    ds = gdal.OpenEx('tmp/cache/AMSR_E_L2_Ocean_B01_200206182340_A.hdf', gdal.OF_MULTIDIM_RASTER, open_options = ['LIST_SDS=YES'])
    rg = ds.GetRootGroup()
    assert rg.GetGroupNames() == ['swaths', 'scientific_datasets']

###############################################################################
# Test reading HDFEOS GRID products


def test_hdf4multidim_hdfeos_grid():

    if not gdaltest.download_file('http://download.osgeo.org/gdal/data/hdf4/MOD09A1.A2010041.h06v03.005.2010051001103.hdf', 'MOD09A1.A2010041.h06v03.005.2010051001103.hdf'):
        pytest.skip()

    ds = gdal.OpenEx('tmp/cache/MOD09A1.A2010041.h06v03.005.2010051001103.hdf', gdal.OF_MULTIDIM_RASTER)
    assert ds
    rg = ds.GetRootGroup()
    assert rg.GetGroupNames() == ['eos_grids']
    attrs = rg.GetAttributes()
    assert attrs
    attr = rg.GetAttribute('HDFEOSVersion')
    assert attr.Read() == 'HDFEOS_V2.9'
    eos_grids = rg.OpenGroup('eos_grids')
    assert eos_grids
    assert not rg.OpenGroup('foo')
    assert eos_grids.GetGroupNames() == ['MOD_Grid_500m_Surface_Reflectance']
    MOD_Grid_500m_Surface_Reflectance = eos_grids.OpenGroup('MOD_Grid_500m_Surface_Reflectance')
    assert MOD_Grid_500m_Surface_Reflectance
    assert not eos_grids.OpenGroup('foo')
    attrs = MOD_Grid_500m_Surface_Reflectance.GetAttributes()
    assert len(attrs) == 0
    assert MOD_Grid_500m_Surface_Reflectance.GetGroupNames() == ['Data Fields']
    assert not MOD_Grid_500m_Surface_Reflectance.OpenGroup('foo')
    datafields = MOD_Grid_500m_Surface_Reflectance.OpenGroup('Data Fields')
    assert datafields
    assert len(datafields.GetMDArrayNames()) == 13
    assert not datafields.OpenMDArray('foo')
    array = datafields.OpenMDArray('sur_refl_b01')
    assert array
    dims = array.GetDimensions()
    assert len(dims) == 2
    assert dims[0].GetFullName() == '/eos_grids/MOD_Grid_500m_Surface_Reflectance/YDim'
    assert dims[0].GetSize() == 2400
    assert dims[1].GetFullName() == '/eos_grids/MOD_Grid_500m_Surface_Reflectance/XDim'
    assert dims[1].GetSize() == 2400
    assert array.GetDataType().GetNumericDataType() == gdal.GDT_Int16
    assert array.GetNoDataValueAsDouble() == -28672.0
    assert array.GetOffset() == 0
    assert array.GetScale() == 0.0001
    assert array.GetUnit() == 'reflectance'
    attr = array.GetAttribute('valid_range')
    assert attr.Read() == (-100, 16000)
    assert array.GetSpatialRef()

    got_data = array.Read(array_start_idx = [2398, 2398], count = [2, 2])
    assert len(got_data) == 2 * 2 * 2
    assert struct.unpack('h' * 4, got_data) == (-24, 0,
                                                -15, -22)

    dims = MOD_Grid_500m_Surface_Reflectance.GetDimensions()
    assert len(dims) == 2
    assert dims[0].GetName() == 'YDim'
    assert dims[0].GetIndexingVariable()
    assert dims[1].GetName() == 'XDim'
    assert dims[1].GetIndexingVariable()

    assert MOD_Grid_500m_Surface_Reflectance.GetMDArrayNames() == ['YDim', 'XDim']
    XDim = MOD_Grid_500m_Surface_Reflectance.OpenMDArray('XDim')
    assert XDim
    YDim = MOD_Grid_500m_Surface_Reflectance.OpenMDArray('YDim')
    assert YDim
    assert not MOD_Grid_500m_Surface_Reflectance.OpenMDArray('foo')

    ds = gdal.OpenEx('tmp/cache/MOD09A1.A2010041.h06v03.005.2010051001103.hdf', gdal.OF_MULTIDIM_RASTER, open_options = ['LIST_SDS=YES'])
    rg = ds.GetRootGroup()
    assert rg.GetGroupNames() == ['eos_grids', 'scientific_datasets']


###############################################################################
# Test reading GDAL SDS 2D


def test_hdf4multidim_gdal_sds_2d():

    ds = gdal.OpenEx('data/byte_2.hdf', gdal.OF_MULTIDIM_RASTER)
    assert ds
    rg = ds.GetRootGroup()
    assert rg.GetGroupNames() is None
    assert rg.OpenGroup('scientific_datasets') is None
    assert rg.GetMDArrayNames() == ['Band0', 'X', 'Y']
    dims = rg.GetDimensions()
    assert len(dims) == 2
    array = rg.OpenMDArray('Band0')
    assert array
    dims = array.GetDimensions()
    assert len(dims) == 2
    assert dims[0].GetFullName() == '/Y'
    assert dims[1].GetFullName() == '/X'
    assert array.GetDataType().GetNumericDataType() == gdal.GDT_Byte

    got_data = array.Read(array_start_idx = [0, 0], count = [2, 2])
    assert len(got_data) == 2 * 2
    assert struct.unpack('B' * 4, got_data) == (107, 123,
                                                115, 132)
    assert array.GetSpatialRef()
    assert array.GetUnit() == ''
    assert not array.GetOffset()
    assert not array.GetScale()
    assert not array.GetNoDataValueAsDouble()

    X = dims[0].GetIndexingVariable()
    assert X
    assert struct.unpack('d' * 2, X.Read(array_start_idx = [0], count = [2])) == (3751290.0, 3751230.0)

    Y = dims[1].GetIndexingVariable()
    assert Y
    assert struct.unpack('d' * 2, Y.Read(array_start_idx = [0], count = [2])) == (440750.0, 440810.0)


###############################################################################
# Test reading GDAL SDS 3D


def test_hdf4multidim_gdal_sds_3d():

    ds = gdal.OpenEx('data/byte_3.hdf', gdal.OF_MULTIDIM_RASTER)
    assert ds
    rg = ds.GetRootGroup()
    assert rg.GetGroupNames() is None
    assert rg.OpenGroup('scientific_datasets') is None
    assert rg.GetMDArrayNames() == ['3-dimensional Scientific Dataset', 'X', 'Y']
    dims = rg.GetDimensions()
    assert len(dims) == 3
    array = rg.OpenMDArray('3-dimensional Scientific Dataset')
    assert array
    dims = array.GetDimensions()
    assert len(dims) == 3
    assert dims[0].GetFullName() == '/Y'
    assert dims[1].GetFullName() == '/X'
    assert dims[2].GetFullName() == '/Band'
    assert array.GetDataType().GetNumericDataType() == gdal.GDT_Byte

    got_data = array.Read(array_start_idx = [0, 0, 0], count = [2, 2, 1])
    assert len(got_data) == 2 * 2
    assert struct.unpack('B' * 4, got_data) == (107, 123,
                                                115, 132)

    got_data = array.Transpose([2,1,0]).Read(array_start_idx = [0, 0, 0], count = [1, 2, 2])
    assert len(got_data) == 2 * 2
    assert struct.unpack('B' * 4, got_data) == (107, 115,
                                                123, 132)

    assert array.GetSpatialRef()

    X = dims[0].GetIndexingVariable()
    assert X
    assert struct.unpack('d' * 2, X.Read(array_start_idx = [0], count = [2])) == (3751290.0, 3751230.0)
    assert rg.OpenMDArray('X')

    Y = dims[1].GetIndexingVariable()
    assert Y
    assert struct.unpack('d' * 2, Y.Read(array_start_idx = [0], count = [2])) == (440750.0, 440810.0)
    assert rg.OpenMDArray('Y')


###############################################################################
# Test reading a simple SDS product


def test_hdf4multidim_sds():

    # Generated with
    # https://support.hdfgroup.org/ftp/HDF/HDF_Current/src/unpacked/mfhdf/examples/SD_create_sds.c
    # + https://support.hdfgroup.org/ftp/HDF/HDF_Current/src/unpacked/mfhdf/examples/SD_set_get_dim_info.c
    # + https://support.hdfgroup.org/ftp/HDF/HDF_Current/src/unpacked/mfhdf/examples/SD_set_attr.c
    ds = gdal.OpenEx('data/SDS.hdf', gdal.OF_MULTIDIM_RASTER)
    assert ds
    rg = ds.GetRootGroup()
    assert rg.GetGroupNames() == ['scientific_datasets']
    subg = rg.OpenGroup('scientific_datasets')
    assert subg
    attrs = rg.GetAttributes()
    assert len(attrs) == 1
    assert attrs[0].GetName() == 'File_contents'
    assert attrs[0].Read() == 'Storm_track_data'
    assert not rg.OpenGroup('foo')
    assert not rg.GetMDArrayNames()
    assert not rg.OpenMDArray('foo')
    assert not subg.GetGroupNames()
    assert not subg.OpenGroup('foo')
    assert subg.GetMDArrayNames() == ['SDStemplate', 'Y_Axis', 'X_Axis']
    dims = subg.GetDimensions()
    assert len(dims) == 2
    array = subg.OpenMDArray('SDStemplate')
    assert array
    dims = array.GetDimensions()
    assert len(dims) == 2
    assert dims[0].GetFullName() == '/scientific_datasets/Y_Axis'
    assert dims[0].GetIndexingVariable()
    assert dims[1].GetFullName() == '/scientific_datasets/X_Axis'
    assert dims[1].GetIndexingVariable()
    attrs = array.GetAttributes()
    assert len(attrs) == 1
    assert attrs[0].GetName() == 'Valid_range'
    attr = array.GetAttribute('Valid_range')
    assert attr
    assert attr.Read() == (2, 10)
    assert array.GetUnit() == ''
    assert not array.GetSpatialRef()


###############################################################################
# Test reading a SDS product with unlimited dimension


def test_hdf4multidim_sds_unlimited_dim():

    # Generated with
    # hhttps://support.hdfgroup.org/ftp/HDF/HDF_Current/src/unpacked/mfhdf/examples/SD_unlimited_sds.c
    ds = gdal.OpenEx('data/SDSUNLIMITED.hdf', gdal.OF_MULTIDIM_RASTER)
    assert ds
    rg = ds.GetRootGroup()
    assert rg.GetGroupNames() == ['scientific_datasets']
    subg = rg.OpenGroup('scientific_datasets')
    assert subg
    dims = subg.GetDimensions()
    assert len(dims) == 2
    assert dims[0].GetName() == 'fakeDim0'
    assert dims[0].GetSize() == 11
    array = subg.OpenMDArray('AppendableData')
    dims = array.GetDimensions()
    assert len(dims) == 2
    assert dims[0].GetName() == 'fakeDim0'
    assert dims[0].GetSize() == 11
    assert len(array.Read()) == 11 * 10 * 4


###############################################################################
# Test reading a 'random' SDS product


def test_hdf4multidim_sds_read_world():

    if not gdaltest.download_file('http://download.osgeo.org/gdal/data/hdf4/A2004259075000.L2_LAC_SST.hdf', 'A2004259075000.L2_LAC_SST.hdf'):
        pytest.skip()

    ds = gdal.OpenEx('tmp/cache/A2004259075000.L2_LAC_SST.hdf', gdal.OF_MULTIDIM_RASTER)
    assert ds
    rg = ds.GetRootGroup()
    assert rg.GetGroupNames() == ['scientific_datasets']
    subg = rg.OpenGroup('scientific_datasets')
    assert subg
    assert not rg.OpenGroup('foo')
    assert not rg.GetMDArrayNames()
    assert not rg.OpenMDArray('foo')
    assert not subg.GetGroupNames()
    assert not subg.OpenGroup('foo')
    assert subg.GetMDArrayNames() == ['sst']
    dims = subg.GetDimensions()
    assert len(dims) == 2
    sst = subg.OpenMDArray('sst')
    assert sst
    dims = sst.GetDimensions()
    assert len(dims) == 2
    assert dims[0].GetFullName() == '/scientific_datasets/fakeDim0'
    assert dims[1].GetFullName() == '/scientific_datasets/fakeDim1'
    attrs = sst.GetAttributes()
    assert len(attrs) == 5
    attr = sst.GetAttribute('long_name')
    assert attr
    assert attr.Read() == 'Sea Surface Temperature'
    assert sst.GetUnit() == 'degrees-C'
    assert not sst.GetSpatialRef()

###############################################################################
# Test reading a SDS product with indexed dimensions


def test_hdf4multidim_sds_read_world_with_indexing_variable():

    if not gdaltest.download_file('https://download.osgeo.org/gdal/data/hdf4/REANALYSIS_1999217.hdf', 'REANALYSIS_1999217.hdf'):
        pytest.skip()

    ds = gdal.OpenEx('tmp/cache/REANALYSIS_1999217.hdf', gdal.OF_MULTIDIM_RASTER)
    assert ds
    rg = ds.GetRootGroup()
    assert rg.GetGroupNames() == ['scientific_datasets']
    subg = rg.OpenGroup('scientific_datasets')
    assert subg
    dims = subg.GetDimensions()
    assert len(dims) == 3
    assert dims[0].GetFullName() == '/scientific_datasets/lat'
    assert dims[0].GetIndexingVariable()
    assert dims[1].GetFullName() == '/scientific_datasets/lon'
    assert dims[1].GetIndexingVariable()
    assert dims[2].GetFullName() == '/scientific_datasets/time'
    assert not dims[2].GetIndexingVariable()
    assert subg.GetMDArrayNames() == ['lat', 'lon', 'slp', 'pr_wtr', 'air']
    slp = subg.OpenMDArray('slp')
    assert slp
    assert slp.GetUnit() == 'Pascals'
    assert slp.GetOffset() == 119765.0
    assert slp.GetScale() == 1.0

###############################################################################
# Test reading a GR dataset


def test_hdf4multidim_gr():

    # Generated by https://support.hdfgroup.org/ftp/HDF/HDF_Current/src/unpacked/hdf/examples/GR_create_and_write_image.c

    ds = gdal.OpenEx('data/General_RImages.hdf', gdal.OF_MULTIDIM_RASTER)
    assert ds
    rg = ds.GetRootGroup()
    assert rg.GetGroupNames() == ['general_rasters']
    subg = rg.OpenGroup('general_rasters')
    assert subg
    attrs = subg.GetAttributes()
    assert len(attrs) == 2
    attr = subg.GetAttribute('File Attribute 1')
    assert attr.Read() == 'Contents of First FILE Attribute'
    assert not subg.GetGroupNames()
    assert not subg.OpenGroup('foo')
    assert subg.GetMDArrayNames() == ['Image Array 1']
    array = subg.OpenMDArray('Image Array 1')
    assert not subg.OpenMDArray('foo')
    dims = array.GetDimensions()
    assert len(dims) == 3
    assert dims[0].GetName() == 'y'
    assert dims[0].GetSize() == 10
    assert dims[1].GetName() == 'x'
    assert dims[1].GetSize() == 5
    assert dims[2].GetName() == 'bands'
    assert dims[2].GetSize() == 2

    got_data = array.Read(array_start_idx = [1, 2, 0], count = [3, 2, 2])
    assert len(got_data) == 3 * 2 * 2 * 2
    assert struct.unpack('h' * 12, got_data) == (4, 5,
                                                 5, 6,
                                                 6, 7,
                                                 5, 6,
                                                 6, 7,
                                                 7, 8)

    got_data = array.Read(array_start_idx = [1, 2, 0], count = [3, 2, 1])
    assert len(got_data) == 3 * 2 * 1 * 2
    assert struct.unpack('h' * 6, got_data) == (4,
                                                5,
                                                6,
                                                5,
                                                6,
                                                7)

    got_data = array.Read(array_start_idx = [1, 2, 1], count = [3, 2, 1])
    assert len(got_data) == 3 * 2 * 1 * 2
    assert struct.unpack('h' * 6, got_data) == (5,
                                                6,
                                                7,
                                                6,
                                                7,
                                                8)

    got_data = array.Read(array_start_idx = [1, 2, 1], count = [3, 2, 2], array_step = [1, 1, -1])
    assert len(got_data) == 3 * 2 * 2 * 2
    assert struct.unpack('h' * 12, got_data) == (5, 4,
                                                 6, 5,
                                                 7, 6,
                                                 6, 5,
                                                 7, 6,
                                                 8, 7)

    attrs = array.GetAttributes()
    assert len(attrs) == 2
    attr = array.GetAttribute('Image Attribute 1')
    assert attr.Read() == "Contents of IMAGE's First Attribute"
    attr = array.GetAttribute('Image Attribute 2')
    assert attr.Read() == (1, 2, 3, 4, 5, 6)

###############################################################################
# Test reading a GR dataset with a palette


def test_hdf4multidim_gr_palette():

    # Generated by https://support.hdfgroup.org/ftp/HDF/HDF_Current/src/unpacked/hdf/examples/GR_write_palette.c

    def get_lut():
        ds = gdal.OpenEx('data/Image_with_Palette.hdf', gdal.OF_MULTIDIM_RASTER)
        assert ds
        rg = ds.GetRootGroup()
        assert rg.GetGroupNames() == ['general_rasters']
        subg = rg.OpenGroup('general_rasters')
        array = subg.OpenMDArray('Image with Palette')
        assert array
        lut = array.GetAttribute('lut')
        return lut

    lut = get_lut()
    assert lut
    got = lut.Read()
    assert len(got) == 3 * 256
    assert got[0] == 0
    assert got[1] == 1
    assert got[2] == 2
    assert got[255 * 3 + 0] == 255
    assert got[255 * 3 + 1] == 0
    assert got[255 * 3 + 2] == 1

