# -*- coding: utf-8 -*-

# sensormapgeo, Transform remote sensing images between sensor and map geometry.
#
# Copyright (C) 2020-2021
# - Daniel Scheffler (GFZ Potsdam, daniel.scheffler@gfz-potsdam.de)
# - Helmholtz Centre Potsdam - GFZ German Research Centre for Geosciences Potsdam,
#   Germany (https://www.gfz-potsdam.de/)
#
# This software was developed within the context of the EnMAP project supported
# by the DLR Space Administration with funds of the German Federal Ministry of
# Economic Affairs and Energy (on the basis of a decision by the German Bundestag:
# 50 EE 1529) and contributions from DLR, GFZ and OHB System AG.
#
# 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.

"""Module to transform 3D arrays between sensor and map geometry (using band-wise Lon/Lat arrays)."""

from typing import Union, Tuple
import multiprocessing
from concurrent.futures import TimeoutError as Timeout
from warnings import warn
import warnings
import platform

from pebble import ProcessPool, ProcessExpired
import numpy as np
from pyproj import CRS

with warnings.catch_warnings():
    # work around https://github.com/pytroll/pyresample/issues/375
    warnings.simplefilter('ignore', UserWarning)
    from pyresample import AreaDefinition

from py_tools_ds.geo.coord_trafo import transform_any_prj

from .transformer_2d import \
    SensorMapGeometryTransformer, _corner_coords_lonlat_to_extent, \
    _move_extent_to_coordgrid, _get_validated_tgt_res
from .version import __version__ as _libver


class SensorMapGeometryTransformer3D(object):
    def __init__(self,
                 lons: np.ndarray,
                 lats: np.ndarray,
                 resamp_alg: str = 'nearest',
                 radius_of_influence: int = 30,
                 mp_alg: str = 'auto',
                 **opts) -> None:
        """Get an instance of SensorMapGeometryTransformer.

        :param lons:    3D longitude array corresponding to the 3D sensor geometry array
        :param lats:    3D latitude array corresponding to the 3D sensor geometry array

        :Keyword Arguments:  (further documentation here: https://pyresample.readthedocs.io/en/latest/swath.html)
            - resamp_alg:           resampling algorithm ('nearest', 'bilinear', 'gauss', 'custom')
            - radius_of_influence:  <float> Cut off distance in meters (default: 30)
                                    NOTE: keyword is named 'radius' in case of bilinear resampling
            - mp_alg                multiprocessing algorithm
                                    'bands': parallelize over bands using multiprocessing lib
                                    'tiles': parallelize over tiles using OpenMP
                                    'auto': automatically choose the algorithm
            - sigmas:               <list of floats or float> [ONLY 'gauss'] List of sigmas to use for the gauss
                                    weighting of each channel 1 to k, w_k = exp(-dist^2/sigma_k^2). If only one channel
                                    is resampled sigmas is a single float value.
            - neighbours:           <int> [ONLY 'bilinear', 'gauss'] Number of neighbours to consider for each grid
                                    point when searching the closest corner points
            - epsilon:              <float> Allowed uncertainty in meters. Increasing uncertainty reduces execution time
            - weight_funcs:         <list of function objects or function object> [ONLY 'custom'] List of weight
                                    functions f(dist) to use for the weighting of each channel 1 to k. If only one
                                    channel is resampled weight_funcs is a single function object.
            - fill_value:           <int or None> Set undetermined pixels to this value (default: 0).
                                    If fill_value is None a masked array is returned with undetermined pixels masked
            - reduce_data:          <bool> Perform initial coarse reduction of source dataset in order to reduce
                                    execution time
            - nprocs:               <int>, Number of processor cores to be used
            - segments:             <int or None> Number of segments to use when resampling.
                                    If set to None an estimate will be calculated
            - with_uncert:          <bool> [ONLY 'gauss' and 'custom'] Calculate uncertainty estimates
                                    NOTE: resampling function has 3 return values instead of 1: result, stddev, count
        """
        # validation
        if lons.ndim != 3:
            raise ValueError('Expected a 3D longitude array. Received a %dD array.' % lons.ndim)
        if lats.ndim != 3:
            raise ValueError('Expected a 3D latitude array. Received a %dD array.' % lats.ndim)
        if lons.shape != lats.shape:
            raise ValueError((lons.shape, lats.shape), "'lons' and 'lats' are expected to have the same shape.")

        self.lats = lats
        self.lons = lons
        self.resamp_alg = resamp_alg
        self.radius_of_influence = radius_of_influence
        self.opts = opts

        # define number of CPUs to use (but avoid sub-multiprocessing)
        #   -> parallelize either over bands or over image tiles
        #      bands: multiprocessing uses multiprocessing.Pool, implemented in to_map_geometry / to_sensor_geometry
        #      tiles: multiprocessing uses OpenMP implemented in pykdtree which is used by pyresample
        self.opts['nprocs'] = opts.get('nprocs', multiprocessing.cpu_count())
        if self.opts['nprocs'] > multiprocessing.cpu_count():
            self.opts['nprocs'] = multiprocessing.cpu_count()
        self.mp_alg = ('bands' if self.lons.shape[2] >= opts['nprocs'] else 'tiles') if mp_alg == 'auto' else mp_alg

    @staticmethod
    def _to_map_geometry_2D(kwargs_dict: dict
                            ) -> Tuple[np.ndarray, tuple, str, int]:
        SMGT2D = SensorMapGeometryTransformer(lons=kwargs_dict['lons'],
                                              lats=kwargs_dict['lats'],
                                              resamp_alg=kwargs_dict['resamp_alg'],
                                              radius_of_influence=kwargs_dict['radius_of_influence'],
                                              **kwargs_dict['init_opts'])
        data_mapgeo, out_gt, out_prj = SMGT2D.to_map_geometry(data=kwargs_dict['data'],
                                                              area_definition=kwargs_dict['area_definition'])

        return data_mapgeo, out_gt, out_prj, kwargs_dict['band_idx']

    def _get_common_target_extent(self,
                                  tgt_epsg: int,
                                  tgt_coordgrid: Tuple[Tuple, Tuple] = None):
        if tgt_epsg == 4326:
            corner_coords_ll = [[self.lons[0, 0, :].min(), self.lats[0, 0, :].max()],  # common UL_xy
                                [self.lons[0, -1, :].max(), self.lats[0, -1, :].max()],  # common UR_xy
                                [self.lons[-1, 0, :].min(), self.lats[-1, 0, :].min()],  # common LL_xy
                                [self.lons[-1, -1, :].max(), self.lats[-1, -1, :].min()],  # common LR_xy
                                ]
            common_tgt_extent = _corner_coords_lonlat_to_extent(corner_coords_ll, tgt_epsg)
        else:
            # get Lon/Lat corner coordinates of geolayers
            UL_UR_LL_LR_ll = [(self.lons[y, x], self.lats[y, x]) for y, x in [(0, 0), (0, -1), (-1, 0), (-1, -1)]]

            # transform them to target projection
            UL_UR_LL_LR_prj = [transform_any_prj(4326, tgt_epsg, x, y) for x, y in UL_UR_LL_LR_ll]

            # separate X and Y
            X_prj, Y_prj = zip(*UL_UR_LL_LR_prj)

            # 3D geolayers, i.e., the corner coordinates have multiple values for multiple bands
            # -> use the outermost coordinates to be sure all data is included
            X_prj = (X_prj[0].min(), X_prj[1].max(), X_prj[2].min(), X_prj[3].max())
            Y_prj = (Y_prj[0].max(), Y_prj[1].max(), Y_prj[2].min(), Y_prj[3].min())

            # get the extent
            common_tgt_extent = (min(X_prj), min(Y_prj), max(X_prj), max(Y_prj))

        if tgt_coordgrid:
            common_tgt_extent = _move_extent_to_coordgrid(common_tgt_extent, *tgt_coordgrid)

        return common_tgt_extent

    def _get_common_area_definition(self,
                                    data: np.ndarray,
                                    tgt_prj: Union[str, int],
                                    tgt_extent: Tuple[float, float, float, float] = None,
                                    tgt_res: Tuple = None,
                                    tgt_coordgrid: Tuple[Tuple, Tuple] = None
                                    ) -> AreaDefinition:
        # get common target extent
        tgt_epsg = CRS(tgt_prj).to_epsg()
        tgt_extent = tgt_extent or self._get_common_target_extent(tgt_epsg, tgt_coordgrid=tgt_coordgrid)

        SMGT2D = SensorMapGeometryTransformer(lons=self.lons[:, :, 0],  # does NOT affect the computed area definition
                                              lats=self.lats[:, :, 0],  # -> only needed for __init__
                                              resamp_alg=self.resamp_alg,
                                              radius_of_influence=self.radius_of_influence,
                                              **self.opts)
        common_area_definition = SMGT2D.compute_areadefinition_sensor2map(data=data[:, :, 0],
                                                                          tgt_prj=tgt_prj,
                                                                          tgt_extent=tgt_extent,
                                                                          tgt_res=tgt_res)

        return common_area_definition

    def to_map_geometry(self,
                        data: np.ndarray,
                        tgt_prj: Union[str, int],
                        tgt_extent: Tuple[float, float, float, float] = None,
                        tgt_res: Tuple[float, float] = None,
                        tgt_coordgrid: Tuple[Tuple, Tuple] = None,
                        area_definition: AreaDefinition = None
                        ) -> Tuple[np.ndarray, tuple, str]:
        """Transform the input sensor geometry array into map geometry.

        :param data:            3D numpy array (representing sensor geometry) to be warped to map geometry
        :param tgt_prj:         target projection (WKT or 'epsg:1234' or <EPSG_int>)
        :param tgt_extent:      extent coordinates of output map geometry array (LL_x, LL_y, UR_x, UR_y) in the tgt_prj
        :param tgt_res:         target X/Y resolution (e.g., (30, 30))
        :param tgt_coordgrid:   target coordinate grid ((x, x), (y, y)):
                                if given, the output extent is moved to this coordinate grid
        :param area_definition: an instance of pyresample.geometry.AreaDefinition;
                                OVERRIDES tgt_prj, tgt_extent, tgt_res and tgt_coordgrid; saves computation time
        """
        if data.ndim != 3:
            raise ValueError(data.ndim, "'data' must have 3 dimensions.")

        if data.shape != self.lons.shape:
            raise ValueError(data.shape, 'Expected a sensor geometry data array with %d rows, %d columns and %d bands.'
                             % self.lons.shape)

        if not tgt_prj and not area_definition:
            raise ValueError(tgt_prj, 'Target projection must be given if area_definition is not given.')

        if tgt_coordgrid:
            tgt_res = _get_validated_tgt_res(tgt_coordgrid, tgt_res)

        init_opts = self.opts.copy()
        if self.mp_alg == 'bands':
            del init_opts['nprocs']  # avoid sub-multiprocessing

        # get common area_definition
        if not area_definition:
            area_definition = self._get_common_area_definition(data, tgt_prj, tgt_extent, tgt_res, tgt_coordgrid)

        args = [dict(
            resamp_alg=self.resamp_alg,
            radius_of_influence=self.radius_of_influence,
            init_opts=init_opts,
            area_definition=area_definition,
            band_idx=band,
            lons=self.lons[:, :, band],
            lats=self.lats[:, :, band],
            data=data[:, :, band],
        ) for band in range(data.shape[2])]

        if self.opts['nprocs'] > 1 and self.mp_alg == 'bands':
            # NOTE: The pebble ProcessPool directly returns the results when available (works like a generator).
            #       This saves a lot of memory compared with map. We also don't use an initializer to share the input
            #       arrays because this would allocate the memory for the input arrays of all bands at once.
            # NOTE: Use a multiprocessing imap iterator here when the OpenMP is finally fixed in the pyresample side:
            #       with multiprocessing.Pool(self.opts['nprocs']) as pool:
            #           res = [res for res in pool.imap_unordered(self._to_map_geometry_2D, args)]
            #           pool.close()
            #           pool.join()
            #           return res
            try:
                # this may cause a deadlock with the GNU OpenMP build, thus each WORKER has a timeout of 10 seconds
                with ProcessPool() as pool:
                    future = pool.map(self._to_map_geometry_2D, args, timeout=10)
                    result = [i for i in future.result()]

            except (Timeout, ProcessExpired):
                # use mp_alg='tiles' instead which uses OpenMP under the hood
                msg = "Switched multiprocessing algorithm from 'bands' to 'tiles' due to a timeout in 'bands' mode. "
                if platform.system() == 'Linux':
                    msg += "Consider using the LLVM instead of the GNU build of OpenMP to fix this issue (install, " \
                           "e.g., by 'conda install -c conda-forge _openmp_mutex=*=1_llvm'."
                warn(msg, RuntimeWarning)

                init_opts = self.opts.copy()
                for argset in args:
                    argset.update(dict(init_opts=init_opts))

                result = [self._to_map_geometry_2D(argsdict) for argsdict in args]

        else:
            result = [self._to_map_geometry_2D(argsdict) for argsdict in args]

        band_inds = [res[-1] for res in result]
        data_mapgeo = np.dstack([result[band_inds.index(i)][0] for i in range(data.shape[2])])
        out_gt = result[0][1]
        out_prj = result[0][2]

        return data_mapgeo, out_gt, out_prj

    @staticmethod
    def _to_sensor_geometry_2D(kwargs_dict: dict
                               ) -> (np.ndarray, int):
        SMGT2D = SensorMapGeometryTransformer(lons=kwargs_dict['lons'],
                                              lats=kwargs_dict['lats'],
                                              resamp_alg=kwargs_dict['resamp_alg'],
                                              radius_of_influence=kwargs_dict['radius_of_influence'],
                                              **kwargs_dict['init_opts'])
        data_sensorgeo = SMGT2D.to_sensor_geometry(data=kwargs_dict['data'],
                                                   src_prj=kwargs_dict['src_prj'],
                                                   src_extent=kwargs_dict['src_extent'])

        return data_sensorgeo, kwargs_dict['band_idx']

    def to_sensor_geometry(self,
                           data: np.ndarray,
                           src_prj: Union[str, int],
                           src_extent: Tuple[float, float, float, float]
                           ) -> np.ndarray:
        """Transform the input map geometry array into sensor geometry.

        :param data:        3D numpy array (representing map geometry) to be warped to sensor geometry
        :param src_prj:     projection of the input map geometry array (WKT or 'epsg:1234' or <EPSG_int>)
        :param src_extent:  extent coordinates of input map geometry array (LL_x, LL_y, UR_x, UR_y) in the src_prj
        """
        # reject bilinear resampling due to https://git.gfz-potsdam.de/EnMAP/sensormapgeo/-/issues/7
        if self.resamp_alg == 'bilinear':
            warn(f"Bilinear resampling is not available in sensormapgeo {_libver} when transforming map geometry to "
                 f"sensor geometry due to changes in upstream packages. Using 'gauss' instead. "
                 "Note that bilinear resampling works in sensormapgeo<=0.5.0.", RuntimeWarning)
            self.resamp_alg = 'gauss'
            # self.opts.pop('radius', None)
            # self.opts['radius_of_influence'] = 30  # default

        if data.ndim != 3:
            raise ValueError(data.ndim, "'data' must have 3 dimensions.")

        init_opts = self.opts.copy()
        if self.mp_alg == 'bands':
            del init_opts['nprocs']  # avoid sub-multiprocessing

        args = [dict(
            resamp_alg=self.resamp_alg,
            radius_of_influence=self.radius_of_influence,
            init_opts=init_opts,
            src_prj=src_prj,
            src_extent=src_extent,
            band_idx=band,
            lons=self.lons[:, :, band],
            lats=self.lats[:, :, band],
            data=data[:, :, band],

        ) for band in range(data.shape[2])]

        if self.opts['nprocs'] > 1 and self.mp_alg == 'bands':
            # NOTE: See the comments in the to_map_geometry() method.
            with multiprocessing.Pool(self.opts['nprocs']) as pool:
                result = [res for res in pool.imap_unordered(self._to_sensor_geometry_2D, args)]
                pool.close()
                pool.join()
        else:
            result = [self._to_sensor_geometry_2D(argsdict) for argsdict in args]

        band_inds = [res[-1] for res in result]
        data_sensorgeo = np.dstack([result[band_inds.index(i)][0] for i in range(data.shape[2])])

        return data_sensorgeo
