#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ******************************************************************************
#
#  Project:  GDAL utils.auxiliary
#  Purpose:  gdal utility functions
#  Author:   Even Rouault <even.rouault at spatialys.com>
#  Author:   Idan Miara <idan@miara.com>
#
# ******************************************************************************
#  Copyright (c) 2015, Even Rouault <even.rouault at spatialys.com>
#  Copyright (c) 2020, Idan Miara <idan@miara.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 os
from numbers import Real
from typing import Optional, Union, Sequence, Tuple, Dict, Any, Iterator, List

from osgeo import gdal
from osgeo_utils.auxiliary.base import get_extension, is_path_like, PathLikeOrStr, enum_to_str, OptionalBoolStr, \
    is_true, \
    MaybeSequence, T

PathOrDS = Union[PathLikeOrStr, gdal.Dataset]
DataTypeOrStr = Union[str, int]
CreationOptions = Optional[Dict[str, Any]]


def DoesDriverHandleExtension(drv: gdal.Driver, ext: str) -> bool:
    exts = drv.GetMetadataItem(gdal.DMD_EXTENSIONS)
    return exts is not None and exts.lower().find(ext.lower()) >= 0


def GetOutputDriversFor(filename: PathLikeOrStr, is_raster=True) -> List[str]:
    filename = os.fspath(filename)
    drv_list = []
    ext = get_extension(filename)
    if ext.lower() == 'vrt':
        return ['VRT']
    for i in range(gdal.GetDriverCount()):
        drv = gdal.GetDriver(i)
        if (drv.GetMetadataItem(gdal.DCAP_CREATE) is not None or
            drv.GetMetadataItem(gdal.DCAP_CREATECOPY) is not None) and \
            drv.GetMetadataItem(gdal.DCAP_RASTER if is_raster else gdal.DCAP_VECTOR) is not None:
            if ext and DoesDriverHandleExtension(drv, ext):
                drv_list.append(drv.ShortName)
            else:
                prefix = drv.GetMetadataItem(gdal.DMD_CONNECTION_PREFIX)
                if prefix is not None and filename.lower().startswith(prefix.lower()):
                    drv_list.append(drv.ShortName)

    # GMT is registered before netCDF for opening reasons, but we want
    # netCDF to be used by default for output.
    if ext.lower() == 'nc' and not drv_list and \
        drv_list[0].upper() == 'GMT' and drv_list[1].upper() == 'NETCDF':
        drv_list = ['NETCDF', 'GMT']

    return drv_list


def GetOutputDriverFor(filename: PathLikeOrStr, is_raster=True, default_raster_format='GTiff',
                       default_vector_format='ESRI Shapefile') -> str:
    if not filename:
        return 'MEM'
    drv_list = GetOutputDriversFor(filename, is_raster)
    ext = get_extension(filename)
    if not drv_list:
        if not ext:
            return default_raster_format if is_raster else default_vector_format
        else:
            raise Exception("Cannot guess driver for %s" % filename)
    elif len(drv_list) > 1:
        print("Several drivers matching %s extension. Using %s" % (ext if ext else '', drv_list[0]))
    return drv_list[0]


def open_ds(filename_or_ds: MaybeSequence[PathOrDS], *args, **kwargs) -> MaybeSequence[gdal.Dataset]:
    if not isinstance(filename_or_ds, PathOrDS.__args__):
        return [open_ds(f) for f in filename_or_ds]
    ods = OpenDS(filename_or_ds, *args, **kwargs)
    return ods.__enter__()


def get_ovr_count(filename_or_ds: PathOrDS) -> int:
    with OpenDS(filename_or_ds) as ds:
        bnd = ds.GetRasterBand(1)
        return bnd.GetOverviewCount()


def get_pixel_size(filename_or_ds: PathOrDS) -> Tuple[Real, Real]:
    ds = open_ds(filename_or_ds)
    geo_transform = ds.GetGeoTransform(can_return_null=True)
    if geo_transform is not None:
        return geo_transform[1], geo_transform[5]
    else:
        return 1, 1


ListOfTupleTT_OrT = List[Union[T, Tuple[T, T]]]


def get_sizes_factors_resolutions(filename_or_ds: PathOrDS, dim: Optional[int] = 0) -> \
     Tuple[ListOfTupleTT_OrT[int], ListOfTupleTT_OrT[Real], ListOfTupleTT_OrT[Real]]:

    ds = open_ds(filename_or_ds)
    bnd = ds.GetRasterBand(1)
    ovr_count = bnd.GetOverviewCount()
    r0 = get_pixel_size(ds)
    s0 = ds.RasterXSize, ds.RasterYSize
    f0 = 1, 1
    sizes = [s0]
    factors = [f0]
    resolutions = [r0]
    for i_overview in range(ovr_count):
        h_overview = bnd.GetOverview(i_overview)
        if h_overview is not None:
            s = h_overview.XSize, h_overview.YSize
            f = s0[0] / s[0], s0[1] / s[1]
            r = r0[0] * f[0], r0[1] * f[1]
            sizes.append(s)
            factors.append(f)
            resolutions.append(r)
    if dim is not None:
        sizes = [x[dim] for x in sizes]
        factors = [x[dim] for x in factors]
        resolutions = [x[dim] for x in resolutions]
    return sizes, factors, resolutions


def get_best_ovr_by_resolutions(requested_res: Real, resolutions: Sequence[Real]):
    for ovr, res in enumerate(resolutions):
        if res > requested_res:
            return max(0, ovr-1)
    return len(resolutions)-1


def get_ovr_idx(filename_or_ds: PathOrDS,
                ovr_idx: Optional[int] = None,
                ovr_res: Optional[Union[int, float]] = None) -> int:
    """
    returns a non-negative ovr_idx, from given mutually exclusive ovr_idx (index) or ovr_res (resolution)
    ovr_idx == None and ovr_res == None => returns 0
    ovr_idx: int >= 0 => returns the given ovr_idx
    ovr_idx: int < 0 => -1 is the last overview; -2 is the one before the last and so on
    ovr_res: float|int => returns the best suitable overview for a given resolution
             meaning the ovr with the lowest resolution which is higher then the request
    ovr_idx: float = x => same as (ovr_idx=None, ovr_res=x)
    """
    if ovr_res is not None:
        if ovr_idx is not None:
            raise Exception(f'ovr_idx({ovr_idx}) and ovr_res({ovr_res}) are mutually exclusive both were set')
        ovr_idx = float(ovr_res)
    if ovr_idx is None:
        return 0
    if isinstance(ovr_idx, Sequence):
        ovr_idx = ovr_idx[0]  # in case resolution in both axis were given we'll consider only x resolution
    if isinstance(ovr_idx, int):
        if ovr_idx < 0:
            overview_count = get_ovr_count(filename_or_ds)
            ovr_idx = max(0, overview_count + 1 + ovr_idx)
    elif isinstance(ovr_idx, float):
        _sizes, _factors, resolutions = get_sizes_factors_resolutions(filename_or_ds)
        ovr_idx = get_best_ovr_by_resolutions(ovr_idx, resolutions)
    else:
        raise Exception(f'Got an unexpected overview: {ovr_idx}')
    return ovr_idx


class OpenDS:
    __slots__ = ['filename', 'ds', 'args', 'kwargs', 'own', 'silent_fail']

    def __init__(self, filename_or_ds: PathOrDS, silent_fail=False, *args, **kwargs):
        self.ds: Optional[gdal.Dataset] = None
        self.filename: Optional[PathLikeOrStr] = None
        if is_path_like(filename_or_ds):
            self.filename = os.fspath(filename_or_ds)
        else:
            self.ds = filename_or_ds
        self.args = args
        self.kwargs = kwargs
        self.own = False
        self.silent_fail = silent_fail

    def __enter__(self) -> gdal.Dataset:

        if self.ds is None:
            self.ds = self._open_ds(self.filename, *self.args, **self.kwargs)
            if self.ds is None and not self.silent_fail:
                raise IOError('could not open file "{}"'.format(self.filename))
            self.own = True
        return self.ds

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.own:
            self.ds = False

    @staticmethod
    def _open_ds(
        filename: PathLikeOrStr,
        access_mode=gdal.GA_ReadOnly,
        ovr_idx: Optional[Union[int, float]] = None,
        open_options: Optional[Union[Dict[str, str], Sequence[str]]] = None,
        logger=None,
    ):
        if not open_options:
            open_options = dict()
        elif isinstance(open_options, Sequence):
            open_options = {k: v for k, v in (s.split('=', 1) for s in open_options)}
        else:
            open_options = dict(open_options)
        ovr_idx = get_ovr_idx(filename, ovr_idx)
        if ovr_idx > 0:
            open_options["OVERVIEW_LEVEL"] = ovr_idx - 1  # gdal overview 0 is the first overview (after the base layer)
        if logger is not None:
            s = 'opening file: "{}"'.format(filename)
            if open_options:
                s = s + " with options: {}".format(str(open_options))
            logger.debug(s)
        open_options = ["{}={}".format(k, v) for k, v in open_options.items()]

        return gdal.OpenEx(str(filename), access_mode, open_options=open_options)


def get_data_type(data_type: Optional[DataTypeOrStr]):
    if data_type is None:
        return None
    if isinstance(data_type, str):
        return gdal.GetDataTypeByName(data_type)
    else:
        return data_type


def get_raster_bands(ds: gdal.Dataset) -> Iterator[gdal.Band]:
    return (ds.GetRasterBand(i + 1) for i in range(ds.RasterCount))


def get_band_types(filename_or_ds: PathOrDS):
    with OpenDS(filename_or_ds) as ds:
        return [band.DataType for band in get_raster_bands(ds)]


def get_band_minimum(band: gdal.Band):
    ret = band.GetMinimum()
    if ret is None:
        band.ComputeStatistics(0)
    return band.GetMinimum()


def get_raster_band(filename_or_ds: PathOrDS, bnd_index: int = 1, ovr_index: Optional[int] = None):
    with OpenDS(filename_or_ds) as ds:
        bnd = ds.GetRasterBand(bnd_index)
        if ovr_index is not None:
            bnd = bnd.GetOverview(ovr_index)
        return bnd


def get_raster_minimum(filename_or_ds: PathOrDS, bnd_index: Optional[int] = 1):
    with OpenDS(filename_or_ds) as ds:
        if bnd_index is None:
            return min(get_band_minimum(bnd) for bnd in get_raster_bands(ds))
        else:
            bnd = ds.GetRasterBand(bnd_index)
            return get_band_minimum(bnd)


def get_raster_min_max(filename_or_ds: PathOrDS, bnd_index: int = 1):
    with OpenDS(filename_or_ds) as ds:
        bnd = ds.GetRasterBand(bnd_index)
        bnd.ComputeStatistics(0)
        min_val = bnd.GetMinimum()
        max_val = bnd.GetMaximum()
        return min_val, max_val


def get_nodatavalue(filename_or_ds: PathOrDS):
    with OpenDS(filename_or_ds) as ds:
        band = next(get_raster_bands(ds))
        return band.GetNoDataValue()


def unset_nodatavalue(filename_or_ds: PathOrDS):
    with OpenDS(filename_or_ds, access_mode=gdal.GA_Update) as ds:
        for b in get_raster_bands(ds):
            b.DeleteNoDataValue()


def get_metadata_item(filename_or_ds: PathOrDS, key: str, domain: str, default: Any = None):
    key = key.strip()
    domain = domain.strip()
    with OpenDS(filename_or_ds) as ds:
        metadata_item = ds.GetMetadataItem(key, domain)
        return metadata_item if metadata_item is not None else default


def get_image_structure_metadata(filename_or_ds: PathOrDS, key: str, default: Any = None):
    return get_metadata_item(filename_or_ds, key=key, domain="IMAGE_STRUCTURE", default=default)


def get_bigtiff_creation_option_value(big_tiff: OptionalBoolStr):
    return "IF_SAFER" if big_tiff is None \
        else big_tiff if bool(big_tiff) and isinstance(big_tiff, str) \
        else str(is_true(big_tiff))


def get_ext_by_of(of: str):
    ext = enum_to_str(of).lower()
    if ext in ['gtiff', 'cog', 'mem']:
        ext = 'tif'
    return '.' + ext


def get_band_nums(ds: gdal.Dataset, band_nums: MaybeSequence[int] = None):
    if not band_nums:
        band_nums = list(range(1, ds.RasterCount + 1))
    elif isinstance(band_nums, int):
        band_nums = [band_nums]
    return band_nums


def get_bands(filename_or_ds: PathOrDS, band_nums: MaybeSequence[int], ovr_idx: Optional[int] = None) -> List[gdal.Band]:
    ds = open_ds(filename_or_ds)
    band_nums = get_band_nums(ds, band_nums)
    bands = []
    for band_num in band_nums:
        band: gdal.Band = ds.GetRasterBand(band_num)
        if band is None:
            raise Exception(f'Could not get band {band_num} from file {filename_or_ds}')
        if ovr_idx:
            band = band.GetOverview(ovr_idx-1)
            if band is None:
                raise Exception(f'Could not get overview {ovr_idx} from band {band_num} of file {filename_or_ds}')
        bands.append(band)
    return bands


def get_scales_and_offsets(bands: MaybeSequence[gdal.Band]) -> Tuple[bool, MaybeSequence[Real], MaybeSequence[Real]]:
    single_band = not isinstance(bands, Sequence)
    if single_band:
        bands = [bands]
    scales = [bnd.GetScale() or 1 for bnd in bands]
    offsets = [bnd.GetOffset() or 0 for bnd in bands]
    is_scaled = any(scale != 1 for scale in scales) or any(offset != 0 for offset in offsets)
    if single_band:
        scales, offsets = scales[0], offsets[0]
    return is_scaled, scales, offsets