Module delta.extensions.sources.tiff

# Copyright © 2020, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration.
# All rights reserved.
#
# The DELTA (Deep Earth Learning, Tools, and Analysis) platform is
# licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0.
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Block-aligned reading from multiple Geotiff files.
"""
import os
import math

import numpy as np
from osgeo import gdal

from delta.imagery import delta_image, rectangle
from delta.extensions.sources.npy import NumpyImage


# Suppress GDAL warnings, errors become exceptions so we get them
gdal.SetConfigOption('CPL_LOG', '/dev/null')
gdal.UseExceptions()

_GDAL_TO_NUMPY_TYPES = {
    gdal.GDT_Byte:    np.dtype(np.uint8),
    gdal.GDT_UInt16:  np.dtype(np.uint16),
    gdal.GDT_UInt32:  np.dtype(np.uint32),
    gdal.GDT_Float32: np.dtype(np.float32),
    gdal.GDT_Float64: np.dtype(np.float64)
}
_NUMPY_TO_GDAL_TYPES = {v: k for k, v in _GDAL_TO_NUMPY_TYPES.items()}

class TiffImage(delta_image.DeltaImage):
    """Images supported by GDAL."""

    def __init__(self, path, nodata_value=None):
        """
        Opens a geotiff for reading.

        Parameters
        ----------
        paths: str or List[str]
            Either a single filename or a list.
            For a list, the images are opened in order as a multi-band image, assumed to overlap.
        nodata_value: dtype of image
            Value representing no data.
        """
        paths = self._prep(path)

        self._path = path
        self._paths = paths
        self._handles = []
        for p in paths:
            if not os.path.exists(p):
                raise Exception('Image file does not exist: ' + p)
            result = gdal.Open(p)
            if result is None:
                raise Exception('Failed to open tiff file %s.' % (p))
            self._handles.append(result)
        self._band_map = []
        for i, h in enumerate(self._handles):
            if h.RasterXSize != self._handles[0].RasterXSize or h.RasterYSize != self._handles[0].RasterYSize:
                raise Exception('Images %s and %s have different sizes!' % (self._paths[0], self._paths[i]))
            for j in range(h.RasterCount):
                self._band_map.append((i, j + 1)) # gdal uses 1-based band indexing

        if nodata_value is None:
            temp = self._gdal_band(0).GetNoDataValue()
            super().__init__(temp)
        else:
            super().__init__(nodata_value)

    def __del__(self):
        self.close()

    def _prep(self, paths): #pylint:disable=no-self-use
        """
        Prepare the file to be opened by other tools (unpack, etc).

        This can be overwritten by subclasses to, for example,
        unpack a zip file to a cache directory.

        Parameters
        ----------
        paths: str or List[str]
            Paths passed to constructor

        Returns
        -------
        Returns a list of underlying files to load instead of the original paths.
        """
        if isinstance(paths, str):
            return [paths]
        return paths

    def __asert_open(self):
        if self._handles is None:
            raise IOError('Operating on an image that has been closed.')

    def close(self):
        """
        Close the image.
        """
        self._handles = None # gdal doesn't have a close function for some reason
        self._band_map = None
        self._paths = None

    def path(self):
        """
        Returns the paths returned by `_prep`.
        """
        return self._path

    def num_bands(self):
        self.__asert_open()
        return len(self._band_map)

    def size(self):
        self.__asert_open()
        return (self._handles[0].RasterYSize, self._handles[0].RasterXSize)

    def _read(self, roi, bands, buf=None):
        self.__asert_open()

        num_bands = len(bands) if bands else self.num_bands()
        if buf is None:
            buf = np.zeros(shape=(num_bands, roi.height(), roi.width()), dtype=self.dtype())
        else:
            s = buf[0, :, :].shape
            if s != (roi.height(), roi.width()):
                raise IOError('Buffer shape should be (%d, %d) but is (%d, %d)!' %
                              (roi.height(), roi.width(), s[0], s[1]))
        if bands:
            for i, b in enumerate(bands):
                band_handle = self._gdal_band(b)
                band_handle.ReadAsArray(yoff=roi.min_y, xoff=roi.min_x,
                                        win_ysize=roi.height(), win_xsize=roi.width(), buf_obj=buf[i, :, :])
        else:
            cur_band = 0
            for h in self._handles:
                h.ReadAsArray(yoff=roi.min_y, xoff=roi.min_x,
                              ysize=roi.height(), xsize=roi.width(),
                              buf_obj=buf[cur_band:cur_band + h.RasterCount, :, :])
                cur_band += h.RasterCount
        return np.transpose(buf, [1, 2, 0])

    def _gdal_band(self, band):
        (h, b) = self._band_map[band]
        ret = self._handles[h].GetRasterBand(b)
        assert ret
        return ret

    def _gdal_type(self, band=0):
        """
        Returns the GDAL data type of the image.
        """
        self.__asert_open()
        return self._gdal_band(band).DataType

    def dtype(self):
        self.__asert_open()
        dtype = self._gdal_type(0)
        if dtype in _GDAL_TO_NUMPY_TYPES:
            return _GDAL_TO_NUMPY_TYPES[dtype]
        raise Exception('Unrecognized gdal data type: ' + str(dtype))

    def bytes_per_pixel(self, band=0):
        """
        Returns
        -------
        int:
            the number of bytes per pixel
        """
        self.__asert_open()
        return gdal.GetDataTypeSize(self._gdal_type(band)) // 8

    def block_size(self):
        """
        Returns
        -------
        (int, int):
            block height, block width
        """
        self.__asert_open()
        band_handle = self._gdal_band(0)
        block_size = band_handle.GetBlockSize()
        return (block_size[1], block_size[0])

    def metadata(self):
        self.__asert_open()
        data = dict()
        h = self._handles[0]
        data['projection'] = h.GetProjection()
        data['geotransform'] = h.GetGeoTransform()
        data['gcps'] = h.GetGCPs()
        data['gcpproj'] = h.GetGCPProjection()
        data['metadata'] = h.GetMetadata()
        return data

    def block_aligned_roi(self, desired_roi):
        self.__asert_open()
        bounds = rectangle.Rectangle(0, 0, width=self.width(), height=self.height())
        if not bounds.contains_rect(desired_roi):
            raise Exception('desired_roi ' + str(desired_roi)
                            + ' is outside the bounds of image with size' + str(self.size()))

        block_height, block_width = self.block_size()
        start_block_x = int(math.floor(desired_roi.min_x     / block_width))
        start_block_y = int(math.floor(desired_roi.min_y     / block_height))
        # Rect max is exclusive
        stop_block_x = int(math.floor((desired_roi.max_x-1) / block_width))
        # The stops are inclusive
        stop_block_y = int(math.floor((desired_roi.max_y-1) / block_height))

        start_x = start_block_x * block_width
        start_y = start_block_y * block_height
        w = (stop_block_x - start_block_x + 1) * block_width
        h = (stop_block_y - start_block_y + 1) * block_height

        # Restrict the output region to the bounding box of the image.
        # - Needed to handle images with partial tiles at the boundaries.
        ans = rectangle.Rectangle(start_x, start_y, width=w, height=h)
        bounds = rectangle.Rectangle(0, 0, width=self.width(), height=self.height())
        return ans.get_intersection(bounds)

    def save(self, path, tile_size=None, nodata_value=None, show_progress=False):
        """
        Save to file, with preprocessing applied.

        Parameters
        ----------
        path: str
            Filename to save to.
        tile_size: (int, int)
            If specified, overwrite block size
        nodata_value: image dtype
            If specified, overwrite nodata value
        show_progress: bool
            Write progress bar to stdout
        """

        write_tiff(path, image=self, nodata=nodata_value, block_size=tile_size,
                   show_progress=show_progress)

def _numpy_dtype_to_gdal_type(dtype): #pylint: disable=R0911
    if dtype in _NUMPY_TO_GDAL_TYPES:
        return _NUMPY_TO_GDAL_TYPES[dtype]
    raise Exception('Unrecognized numpy data type: ' + str(dtype))

def write_tiff(output_path: str, data: np.ndarray=None, image: delta_image.DeltaImage=None,
               nodata=None, metadata: dict=None, block_size=None, show_progress: bool=False):
    """
    Write a numpy array to a file as a tiff.

    Parameters
    ----------
    output_path: str
        Filename to save tiff file to
    data: numpy.ndarray
        Image data to save.
    image: delta_image.DeltaImage
        Image data to save (specify one of this or data).
    nodata: Any
        Nodata value.
    metadata: dict
        Optional metadata to include.
    block_size: Tuple[int]
        Optionally override block size for writing.
    show_progress: bool
        Display command line progress bar.
    """

    if data is not None:
        assert image is None, 'Must specify one of data or image.'
        image = NumpyImage(data, nodata_value=nodata)
    assert image is not None, 'Must specify either data or image.'
    num_bands = image.num_bands()
    data_type = _numpy_dtype_to_gdal_type(image.dtype())
    size = image.size()
    # gdal requires block size of 16 when writing...
    ts = image.block_size()
    if block_size:
        ts = block_size
    if nodata is None:
        nodata = image.nodata_value()
    if metadata is None:
        metadata = image.metadata()

    with _TiffWriter(output_path, height=size[0], width=size[1], num_bands=num_bands,
                     data_type=data_type, metadata=metadata, nodata_value=nodata,
                     tile_height=ts[0], tile_width=ts[1]) as writer:
        input_bounds = rectangle.Rectangle(0, 0, width=size[1], height=size[0])
        ts = writer.tile_shape()
        output_rois = input_bounds.make_tile_rois_yx(ts, include_partials=True)[0]
        def callback_function(output_roi, data, _):
            """Callback function to write the first channel to the output file."""

            # Figure out some ROI positioning values
            block_x = output_roi.min_x // ts[1]
            block_y = output_roi.min_y // ts[0]

            # Loop on bands
            if len(data.shape) == 2:
                writer.write_block(data[:, :], block_y, block_x, 0)
            else:
                for band in range(num_bands):
                    writer.write_block(data[:, :, band], block_y, block_x, band)

        image.process_rois(output_rois, callback_function, show_progress=show_progress)

class _TiffWriter:
    """
    Class to manage block writes to a Geotiff file. Internal helper class.
    """
    def __init__(self, path, height, width, num_bands=1, data_type=gdal.GDT_Byte, #pylint:disable=too-many-arguments
                 tile_height=256, tile_width=256, nodata_value=None, metadata=None):
        self._width  = width
        self._height = height
        self._tile_height = tile_height
        self._tile_width  = tile_width
        self._handle = None

        self.__initialize(path, num_bands, data_type, nodata_value, metadata)

    def __initialize(self, path, num_bands, data_type, nodata_value, metadata):
        options = ['BigTIFF=IF_SAFER', 'INTERLEAVE=BAND']
        if data_type not in (gdal.GDT_Float32, gdal.GDT_Float64):
            options += ['COMPRESS=LZW']

        MIN_SIZE_FOR_TILES=100
        if self._width > MIN_SIZE_FOR_TILES or self._height > MIN_SIZE_FOR_TILES:
            options += ['TILED=YES']
            # requires 16 byte alignment in tiled mode
            self._tile_width = min(max(1, self._tile_width // 16) * 16, 1024)
            self._tile_height = min(max(1, self._tile_height // 16) * 16, 1024)

        options += ['BLOCKXSIZE='+str(self._tile_width),
                    'BLOCKYSIZE='+str(self._tile_height)]

        driver = gdal.GetDriverByName('GTiff')
        self._handle = driver.Create(path, ysize=self._height, xsize=self._width,
                                     bands=num_bands, eType=data_type, options=options)
        if not self._handle:
            raise Exception('Failed to create output file: ' + path)

        if nodata_value is not None:
            for i in range(1,num_bands+1):
                self._handle.GetRasterBand(i).SetNoDataValue(nodata_value)

        if metadata:
            self._handle.SetProjection  (metadata['projection'  ])
            self._handle.SetGeoTransform(metadata['geotransform'])
            self._handle.SetMetadata    (metadata['metadata'    ])
            self._handle.SetGCPs        (metadata['gcps'], metadata['gcpproj'])

    def __del__(self):
        self.close()

    def __enter__(self):
        return self

    def __exit__(self, *unused):
        self.close()
        return False

    def tile_shape(self):
        return (self._tile_height, self._tile_width)

    def close(self):
        if self._handle is not None:
            self._handle.FlushCache()
            self._handle = None

    def get_num_tiles(self):
        num_x = int(math.ceil(self._width  / self._tile_width))
        num_y = int(math.ceil(self._height / self._tile_height))
        return (num_y, num_x)

    def write_block(self, data, block_y, block_x, band=0):
        '''Add a tile write command to the queue.
           Partial tiles are allowed at the right at bottom edges.
        '''

        # Check that the tile position is valid
        num_tiles = self.get_num_tiles()
        if (block_x >= num_tiles[1]) or (block_y >= num_tiles[0]):
            raise Exception('Block position ' + str((block_x, block_y))
                            + ' is outside the tile count: ' + str(num_tiles))
        is_edge_block = ((block_x == num_tiles[1]-1) or
                         (block_y == num_tiles[0]-1))

        if is_edge_block: # Data must fit inside the image size
            max_x = block_x * self._tile_width  + data.shape[1]
            max_y = block_y * self._tile_height + data.shape[0]
            if max_x > self._width or max_y > self._height:
                raise Exception('Error: Data block max position '
                                + str((max_y, max_x))
                                + ' falls outside the image bounds: '
                                + str((self._height, self._width)))
        else: # Shape must be exactly one tile
            if ( (data.shape[1] != self._tile_width) or
                 (data.shape[0] != self._tile_height)  ):
                raise Exception('Error: Data block size is ' + str(data.shape)
                                + ', output file block size is '
                                + str((self._tile_height, self._tile_width)))

        gdal_band = self._handle.GetRasterBand(band+1)
        assert gdal_band

        gdal_band.WriteArray(data, yoff=block_y * self._tile_height,
                             xoff=block_x * self._tile_width)

    def write_region(self, data, y, x):
        assert 0 <= y < self._height
        assert 0 <= x < self._width

        if len(data.shape) < 3:
            gdal_band = self._handle.GetRasterBand(1)
            assert gdal_band
            gdal_band.WriteArray(data, yoff=y, xoff=x)
            return

        for band in range(data.shape[2]):
            gdal_band = self._handle.GetRasterBand(band+1)
            assert gdal_band
            gdal_band.WriteArray(data[:, :, band], yoff=y, xoff=x)

class TiffWriter(delta_image.DeltaImageWriter):
    """
    Write a geotiff to a file.
    """
    def __init__(self, filename):
        self._filename = filename
        self._tiff_w = None

    def initialize(self, size, numpy_dtype, metadata=None, nodata_value=None):
        assert (len(size) == 3), ('Error: len(size) of '+str(size)+' != 3')
        TILE_SIZE = 256
        self._tiff_w = _TiffWriter(self._filename, size[0], size[1], num_bands=size[2],
                                   data_type=_numpy_dtype_to_gdal_type(numpy_dtype), metadata=metadata,
                                   nodata_value=nodata_value,
                                   tile_height=min(TILE_SIZE, size[0]), tile_width=min(TILE_SIZE, size[1]))

    def write(self, data, y, x):
        self._tiff_w.write_region(data, y, x)

    def close(self):
        if self._tiff_w is not None:
            self._tiff_w.close()

    def abort(self):
        self.close()
        try:
            os.remove(self._filename)
        except OSError:
            pass