# vim: ts=4:sts=4:sw=4
#
# @date 2019-09-25
#
# This file is part of IDF converter, a set of tools to convert satellite,
# in-situ and numerical model data into Intermediary Data Format, making them
# compatible with the SEAScope application.
#
# Copyright (C) 2014-2022 OceanDataLab
#
# IDF converter is free software: you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published by the
# Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.
#
# IDF converter is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with IDF converter. If not, see <https://www.gnu.org/licenses/>.

"""This module contains methods for removing signal portions generated by the
traversal of atmosphere."""

import numpy
import typing
import logging
import scipy.interpolate
from idf_converter.lib.types import AtmCorrData, FormatterResult
from idf_converter.lib.types import InputOptions, OutputOptions, Granule

logger = logging.getLogger(__name__)


class NonFiniteValueFound(ValueError):
    """Error raised when an infinite or a NaN value is present in the
    atmospheric correction, most likely due to invalid values in the lookup
    table."""
    def __init__(self) -> None:
        """"""
        _msg = ('Infinite or nan value found in the atmospheric correction, '
                'please check that the LUT has valid values for the angles '
                'contained in the input data file')
        super(NonFiniteValueFound, self).__init__(_msg)


class IncompatibleLUTSize(ValueError):
    """Error raised when the number of entries in the LUT is not a multiple of
    (number of observation zenith angles * number of zenithal deltas)."""
    def __init__(self, lut_size: int, lut_size_factor: float) -> None:
        """"""
        self.lut_size = lut_size
        self.lut_size_factor = lut_size_factor
        _msg = ('The size of the LUT does not match the hardcoded '
                'parameters: the number of lines per sza value should '
                f'be {lut_size_factor}')
        super(IncompatibleLUTSize, self).__init__(_msg)


def s3_cnes_rayleigh(input_opts: InputOptions,
                     output_opts: OutputOptions,
                     granule: Granule,
                     *args: typing.Iterable[typing.Any],
                     targets: typing.Iterable[str],
                     atm_corr_data: AtmCorrData) -> FormatterResult:
    """Apply correction for Rayleigh scattering based on lookup tables
    generated by CNES.

    Parameters
    ----------
    input_opts: dict
        Options and information related to input data
    output_opts: dict
        Options and information related to formatting and serialization to IDF
        format
    granule: idf_converter.lib.Granule
        Object representing the granule data and metadata
    targets: tuple
        Identifers for variables on which the mask will be applied
    atm_corr_data: dict
        Dictionary containing all the variables and parameters required for
        computing the correction

    Returns
    -------
    tuple(dict, dict, idf_converter.lib.Granule)
        A tuple which contains three elements:

        - the input_options :obj:dict passed to this method

        - the output_options :obj:dict passed to this method

        - the :obj:`idf_converter.lib.Granule` where the Rayleigh scattering
          has been removed from the targetted variables
    """
    # LUT data
    sza_lut = atm_corr_data['SZA_LUT']
    oza_lut = atm_corr_data['OZA_LUT']
    delta_lut = atm_corr_data['DELTA_LUT']
    rho_atm = atm_corr_data['RHO_LUT']

    # Granule angles
    sza = atm_corr_data['SZA']  # Sun Zenithal Angle
    saa = atm_corr_data['SAA']  # Sun Azimuthal Angle
    oaa = atm_corr_data['OAA']  # Observation Azimuthal Angle
    oza = atm_corr_data['OZA']  # Observation Zenithal Angle

    # Extra-terrestrial solar irradiance
    solar_flux = atm_corr_data['solar_flux']

    bands_count = atm_corr_data['bands_count']
    bands_index = atm_corr_data['bands_index']
    ac_subsampling = atm_corr_data['ac_subsampling']
    al_subsampled = atm_corr_data['al_subsampled']
    nrow = atm_corr_data['nrow']
    ncell = atm_corr_data['ncell']

    oza_resol = 5
    delta_resol = 10
    ntot_delta = 360 / delta_resol + 1
    ntot_oza = 90 / oza_resol + 1

    # Check the size of the LUT
    lut_size = len(rho_atm)
    lut_size_factor = ntot_oza * ntot_delta
    if 0 != (lut_size % lut_size_factor):
        raise IncompatibleLUTSize(lut_size, lut_size_factor)

    # Delta = Sun Azimuthal Angle - Observation Azimuthal Angle
    delta = saa - oaa
    delta = (delta + 360) % 360

    # Interpolate atmospheric correction
    nrow_geom = numpy.shape(sza)[0]
    ncell_geom = numpy.shape(sza)[1]

    rho = numpy.zeros((bands_count, nrow, ncell))
    _points = numpy.array((sza.flatten(), oza.flatten(), delta.flatten())
                          ).transpose()
    sza_lut = numpy.unique(numpy.array(sza_lut, dtype='float32'))
    oza_lut = numpy.unique(numpy.array(oza_lut, dtype='float32'))
    delta_lut = numpy.unique(numpy.array(delta_lut, dtype='float32'))
    lut_points = (sza_lut, oza_lut, delta_lut)
    rho_atm = numpy.array(rho_atm[:][:])
    lut_shape = (sza_lut.size, oza_lut.size, delta_lut.size, 1)

    dst_range_ac = numpy.arange(ncell)
    interpolator = scipy.interpolate.RegularGridInterpolator

    al_rows = nrow_geom
    if al_subsampled is True:
        # Avoid issues in case geometry has more rows than expected (might
        # happen in some edge case)
        al_rows = numpy.minimum(nrow_geom,
                                numpy.floor(nrow / 2.0)).astype('int')
    for var_id in targets:
        iband = bands_index[var_id]

        lut_values = numpy.array((rho_atm[:, iband])).reshape(lut_shape)
        lut_func = interpolator(lut_points, lut_values, method='linear',
                                bounds_error=False, fill_value=None)

        _values = lut_func(_points)
        _values = _values.reshape(nrow_geom, ncell_geom)
        _values = _values * numpy.cos(sza * numpy.pi / 180.)

        src_range_ac = numpy.arange(_values.shape[1]) * ac_subsampling
        # Compensate for along track subsampling (no interpolation along track)
        for irow in range(0, al_rows):
            src_values = _values[irow]

            if al_subsampled is True:
                # Right now the only example we have is Sentinel-3 SLSTR which
                # has a fixed /2 subsampling.
                rho[iband, irow * 2] = numpy.interp(dst_range_ac, src_range_ac,
                                                    src_values)
                irownext = irow * 2 + 1
                if irownext >= nrow:
                    break
                rho[iband, irownext] = numpy.interp(dst_range_ac, src_range_ac,
                                                    src_values)
            else:
                # TODO: check if the "irow < nrow" check is required here, it
                # was in the original OLCI L2 reader
                if irow >= nrow:
                    break

                rho[iband, irow] = numpy.interp(dst_range_ac, src_range_ac,
                                                src_values)

        # TODO check if there are still nan
        rho[numpy.where(numpy.isnan(rho))] = 0

        # 10**(-3) variation for solar flux is neglectable so we consider the
        # mean
        es = numpy.mean(solar_flux[var_id])

        # Convert reflectance into radiance TOA:
        toa_rad = rho[iband, :, :] * es / numpy.pi

        # An erroneous LUT could produce nan values for the atmospheric
        # correction. This is not acceptable.
        if numpy.any(~numpy.isfinite(toa_rad)):
            raise NonFiniteValueFound()

        # Apply correction
        result = granule.vars[var_id]['array']
        corrected = result - toa_rad
        granule.vars[var_id]['array'] = corrected

    return (input_opts, output_opts, granule)
