# Copyright (c) 2020 CNES/JPL
#
# All rights reserved. Use of this source code is governed by a
# BSD-style license that can be found in the LICENSE file.
"""
Copyright (C) 2017-2021 OceanDataLab
This file is part of skimulator.

skimulator is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

skimulator 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 General Public License for more details.

You should have received a copy of the GNU General Public License
along with skimulator.  If not, see <http://www.gnu.org/licenses/>.
"""

"""
Random signal generation utilities
----------------------------------
"""
from typing import Optional, Tuple
import warnings
try:
    import numba as nb
except ImportError:
    nb = None

import numpy
import scipy.interpolate
try:
    import mkl_fft
    IFFT = mkl_fft.ifft
    IFFT2 = mkl_fft.ifft2
except ImportError:
    IFFT = numpy.fft.ifft
    IFFT2 = numpy.fft.ifft2


def gen_signal_1d(fi: numpy.ndarray,
                  psi: numpy.ndarray,
                  x: numpy.ndarray,
                  nseed: int = 0,
                  fmin: Optional[float] = None,
                  fmax: Optional[float] = None,
                  alpha: int = 10,
                  lf_extpl: bool = False,
                  hf_extpl: bool = False) -> numpy.ndarray:
    """Generate 1d random signal using Fourier coefficient"""
    # Make sure fi, PSi does not contain the zero frequency:
    psi = psi[fi > 0]
    fi = fi[fi > 0]

    # Adjust fmin and fmax to fi bounds if not specified
    fmin = fmin or fi[0]
    fmax = fmax or fi[-1]

    # Go alpha times further in frequency to avoid interpolation aliasing.
    fmaxr = alpha * fmax

    # Interpolation of the non-zero part of the spectrum
    f = numpy.arange(fmin, fmaxr + fmin, fmin)
    mask = psi > 0
    ps = numpy.exp(numpy.interp(numpy.log(f), numpy.log(fi[mask]),
                   numpy.log(psi[mask])))

    # lf_extpl=True prolongates the PSi as a plateau below min(fi).
    # Otherwise, we consider zeros values. same for hf
    ps[f < fi[0]] = psi[0] if lf_extpl else 0
    ps[f > fi[-1]] = psi[-1] if hf_extpl else 0
    ps[f > fmax] = 0

    # Detect the sections (if any) where PSi==0 and apply it to PS
    mask = numpy.interp(f, fi, psi)
    ps[mask == 0] = 0

    f_size = f.size
    phase = numpy.empty((2 * f_size + 1))
    numpy.random.seed(nseed)
    phase[1:(f_size + 1)] = numpy.random.random(f_size) * 2 * numpy.pi
    phase[0] = 0
    phase[-f_size:] = -phase[1:(f_size + 1)][::-1]

    fft1a = numpy.concatenate((numpy.array([0]), 0.5 * ps, 0.5 * ps[::-1]),
    axis=0)
    fft1a = numpy.sqrt(fft1a) * numpy.exp(1j * phase) / fmin**0.5

    yg = 2 * fmaxr * numpy.real(IFFT(fft1a))
    xg = numpy.linspace(0, 0.5 / fmaxr * yg.shape[0], yg.shape[0])

    return numpy.interp(numpy.mod(x, xg.max()), xg, yg)


def calculate_ps2d(f: numpy.ndarray, f2: numpy.ndarray, ps1d: numpy.ndarray,
                   dfx: numpy.ndarray, dfy: numpy.ndarray) -> numpy.ndarray:
    result = numpy.zeros(f2.shape)
    view = result.ravel()
    dfx_2 = dfx * 0.5
    dfx_y = dfx * dfy
    for idx in range(-1, -f.size - 1, -1):
        item = f[idx]
        mask = (f2 >= (item - dfx_2)) & (f2 < (item + dfx_2))
        amount = numpy.sum(result[:, idx]) * dfx_y
        miss = ps1d[idx] * dfx - amount
        view[mask.ravel()] = 0 if miss <= 0 else miss * 0.5 / dfx_y
    return result


def calculate_signal(rectangle, x, y, xgmax, ygmax):
    result = numpy.zeros((len(y), len(x)))
    xn = (x.max() - x[0]) // xgmax
    yn = (y.max() - y[0]) // ygmax
    dx = x - x[0]
    dy = y - y[0]

    for ix_n in range(int(xn + 1)):
        ix0 = numpy.where((dx >= (ix_n * xgmax))
                          & (dx < ((ix_n + 1) * xgmax)))[0]
        for iy_n in range(int(yn + 1)):
            iy0 = numpy.where((dy >= (iy_n * ygmax))
                               & (dy < ((iy_n + 1) * ygmax)))[0]
            result[iy0[0]:iy0[-1] + 1,
                   ix0[0]:ix0[-1] + 1] = rectangle[:len(iy0), :len(ix0)]
    return result


def convert_dbkm2ms(var_out: numpy.array, ac_angle: numpy.array,
                    vsat: float, beam_angle: float)-> numpy.ndarray:
    # Computation of lobe antenna
    th3 = numpy.deg2rad(0.65)
    an = th3 / numpy.sqrt(4 * numpy.log(2))
    # Computation of mispointing
    delta = an**2 / numpy.sin(numpy.deg2rad(beam_angle))**2 * var_out
    out = delta * vsat * numpy.cos(ac_angle)
    return out


def proj_radial(u: numpy.array, v: numpy.array, radial_angle: numpy.array
               ) -> numpy.ndarray:
    ur = u * numpy.cos(radial_angle) + v * numpy.sin(radial_angle)
    return ur


def dist_sphere(lon1, lon2, lat1, lat2, R):
    radlat1 = numpy.deg2rad(lat1)
    radlat2 = numpy.deg2rad(lat2)
    raddifflon = numpy.deg2rad(lon2 - lon1)
    d = numpy.arccos(numpy.sin(radlat2) * numpy.sin(radlat1)
                     + numpy.cos(radlat2) * numpy.cos(radlat1)
                     * numpy.cos(raddifflon))
    return d * R / 1000. # in km


def cross_product(mat: numpy.array, ncoeff: int, cshape: numpy.array
                 )-> numpy.ndarray:
    ncross = sum(range(ncoeff + 1)) + ncoeff
    cross = numpy.full((cshape, ncross), numpy.nan)
    cross[:, 0:ncoeff] = mat[:, :ncoeff]
    niter = ncoeff
    for k in range(ncoeff):
        for l in range(ncoeff - k):
            cross[:, niter] = mat[:, k]*mat[:, k + l]
            niter += 1
    return cross


def reconstruct_var(ncoeff, cshape, f1, b1, cross, xlabel):
    shape_f1 = numpy.shape(f1)[0]
    shape_b1 = numpy.shape(b1)[0]
    #f1 = f1.reshape(int(shape_f1//float(shape_b1)), shape_f1)
    proba = numpy.exp(numpy.dot(cross, f1) + b1)
    proba_tot = numpy.tile(numpy.nansum(proba, axis=1), (shape_b1, 1))
    proba = proba / numpy.transpose(proba_tot)
    var = numpy.sum(proba * numpy.tile(xlabel, (cshape, 1)),axis=1)
    b1 = numpy.tile(b1, (1, 1))
    train_weight = numpy.concatenate((f1,b1),axis=0)
    l_train = numpy.zeros((cshape, ncoeff + 1)).astype('float64')
    l_train[:, : ncoeff] = cross
    l_train[:, ncoeff] = 1.
    res = numpy.dot(l_train, train_weight)
    xx2D = numpy.tile(xlabel, (cshape, 1))
    var = (numpy.sum(xx2D * numpy.exp(res), axis=1)
           / numpy.sum(numpy.exp(res),axis=1))
    return var


if nb is not None:
    _calculate_ps2d = nb.njit("(float64[:, ::1]) (float64[::1],"
                               "float64[:, ::1], float64[::1], float64,"
                               "float64)", cache=True,
                               nogil=True)(calculate_ps2d)
    _calculate_signal = nb.njit("(float64[:, ::1]) (float64[:, ::1],"
                                "float64[::1], float64[::1], float64,"
                                "float64)", cache=True,
                                nogil=True)(calculate_signal)
else:
    _calculate_ps2d = calculate_ps2d
    _calculate_signal = calculate_signal
