"""
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/>.
"""

import numpy
import math
import logging
from typing import Dict, Tuple
import pkg_resources
from . import utils


# Define logger level for debug purposes
logger = logging.getLogger(__name__)


def load_mss_noise() -> Dict[str, numpy.array]:
    mss_path =  pkg_resources.resource_filename('skimulator',
                                                'share/noise_pdf_mss1d.npy')
    mssr_noise = numpy.load(mss_path, allow_pickle=True)[()]
    return mssr_noise


def combine_usr(lon: numpy.array, lat: numpy.array, usr: numpy.array,
                mssx: numpy.array, mssy: numpy.array, mssxy: numpy.array,
                dazi: float, angle: numpy.array, incl: numpy.array,
                wnd_dir: numpy.array, nseed: int=1, Rearth: float=6378. * 10**3
                ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
    '''
    PURPOSE : recombine measurements of various azimuths and ambiguity
    removal using wind direction
    Provides a single estimate (not applicable to a set of points)
    INPUT : - usrazi : radial Stokes drift for each azimuth in azibin [m/s], len = nazi
    dazi : azimuth increment [degrees]
    azimuthat : along track azimuth [radians]
    thwnd : direction of the wind [radians], without ambiguity. Must be in the [0;2*pi[ interval
    OUTPUT : usrhat : estimate of the radial stokes drift [m/s]
    '''
    mssnoise = load_mss_noise()
    #Define ensemble of angle where observations are required
    azibin = numpy.arange(0, 180, dazi)
    # Ambiguity from SKIM measurment
    # TOREMOVE
    usrabs = + numpy.abs(usr)
    #usrabs = + usr # numpy.abs(usr)
    nsample, nbeam = numpy.shape(usrabs)
    # From rad to deg and considering pi ambiguity,
    # Radial angle is refered to the East
    degangle = numpy.rad2deg(numpy.mod(angle, numpy.pi))
    # Inclination of pass (along track angle)
    inclination = numpy.transpose(([incl, ]* nbeam))
    angle_al = numpy.mod(inclination - angle - numpy.pi/2, 2*numpy.pi)
    mssr = (numpy.cos(angle_al)**2 * mssx + numpy.sin(angle_al)**2 * mssy
            + numpy.sin(2*angle_al) * mssxy)
    numpy.random.seed(nseed)
    mssr_noise = numpy.random.choice(mssnoise['bins'], numpy.shape(mssr),
                                     p=mssnoise['hist'])
    mssr = mssr + mssr_noise
    az_al = numpy.mod(inclination - numpy.pi, numpy.pi)
    degangle = numpy.mod(numpy.rad2deg(angle_al), 180)
    az_al = numpy.mod(numpy.rad2deg(az_al), 180)
    # Initialize the reconstruction of ussr
    usr_comb = numpy.full((nsample, nbeam), numpy.nan)
    mssr_comb = numpy.full((nsample, nbeam), numpy.nan)
    usp_comb = numpy.full((nsample, nbeam), numpy.nan)
    # Distance max to take neighbours
    dmax = 200
    for ibeam in range(nbeam):
        for isample in range(nsample):
            if not numpy.isfinite(usrabs[isample, ibeam]):
                continue
            plon = lon[isample, ibeam]
            plat = lat[isample, ibeam]
            pangle = degangle[isample, ibeam]
            rangle = angle_al[isample, ibeam]
            iincl = az_al[isample, ibeam]
            #dist = mod_tools.dist_sphere(plon, lon, plat, lat)
            dist = 110*numpy.sqrt((numpy.cos(numpy.deg2rad(lat)) * (lon - plon))**2 + (lat-plat)**2)
            Idist0 = numpy.where(dist < dmax)
            azimuthr = angle_al[Idist0].ravel()
            lonr = lon[Idist0].ravel()
            latr = lat[Idist0].ravel()
            usrr = usrabs[Idist0].ravel()
            mssrr = mssr[Idist0].ravel()
            distr = dist[Idist0].ravel()
            ang_azi = numpy.digitize(degangle[Idist0].ravel(), azibin) - 1
            Ifinite = numpy.logical_and(numpy.isfinite(usrr),
                                        numpy.isfinite(mssrr))
            # for each bin, find the closest point and append radial Stokes
            # drift with ambiguity on sign
            usrazi = numpy.full(len(azibin), numpy.nan)
            mssrazi = numpy.full(len(azibin), numpy.nan)
            thazi = numpy.full(len(azibin), numpy.nan)
            for iazi in range(len(azibin)):
                I1 = numpy.where(numpy.logical_and(ang_azi==iazi, Ifinite))
                distazi = utils.dist_sphere(plon, lonr[I1], plat, latr[I1], Rearth)
                # TODO: detect if distazi is empty
                try:
                    Idist = numpy.argmin(distazi)
                    usrazi[iazi] = usrr[I1][Idist]
                    mssrazi[iazi] = mssrr[I1][Idist]
                    thazi[iazi] = numpy.mod(azimuthr[I1][Idist], numpy.pi)
                except:
                    usrazi[iazi] = numpy.nan
                    mssrazi[iazi] = numpy.nan
            try:
                iaziat = int(numpy.where(numpy.logical_and(iincl > (azibin),
                                         iincl < (azibin + dazi)))[0])
                Iaziatnot = numpy.logical_not(numpy.arange(len(azibin))==iaziat)
            except:
                iaziat = None
            if iaziat:
            # Interpolate input data in the along track direction (will be to
            # avoid as the noise cnnot be removed in that direction)
                usrazi[iaziat] = numpy.interp(azibin[iaziat],
                                              azibin[Iaziatnot],
                                              usrazi[Iaziatnot], period=180)
                mssrazi[iaziat] = numpy.interp(azibin[iaziat],
                                               azibin[Iaziatnot],
                                               mssrazi[Iaziatnot], period=180)
            # 3. Perform an estimate of the Stokes drift direction and magnitude
            angle_usr = thazi - rangle
            _usr_comb, _usp_comb = proj_uss(usrazi, azibin, angle_usr, dazi,
                                            wnd_dir[isample, ibeam])
            usr_comb[isample, ibeam] = _usr_comb
            usp_comb[isample, ibeam] = _usp_comb
            mssr_comb[isample, ibeam] = numpy.nanmean(mssrazi)
            #mssr_comb[isample, ibeam] = (numpy.nansum(mssrazi * numpy.sin(angle_usr)) * 2 * dazi / 180.)
    return usr_comb, usp_comb, mssr_comb


def proj_uss(usrazi: numpy.array, azibin: numpy.array, angle_usr: numpy.array,
             dazi: float, wnd_dir: numpy.array):
    ''' Project uss on spectrum point
    ARGS:
         usrazi : radial Stokes drift for each azimuth in azibin [m/s]
         azibin: discretized azimuth
         angle_usr: angle of measurement
         dazi: increment angle
         wnd_dir: Wind direction from the ...
    RETURN:
         _usr_comb: Radial component of Stockes drift
         _usp_comb: Perpendicular component of Stockes drift
    '''
    c = numpy.fft.fft(usrazi)/len(usrazi)
    rd2 = -0.5 * numpy.angle(c[1]) # estimate of Us direction
    thrbeam = numpy.mod(rd2, 2 * numpy.pi)
    if numpy.cos(wnd_dir - thrbeam) < 0.:
        thrbeam = numpy.arctan2(-numpy.sin(thrbeam), -numpy.cos(thrbeam))
    # 4. Ambiguity removal using closeness to estimated direction
    radazibin = numpy.deg2rad(azibin)
    cond = numpy.where(numpy.cos(thrbeam - radazibin) > 0)[0]
    usfull = - usrazi * 1
    #TOREMOVE
    #usfull = usrazi
    if cond.any():
        usfull[(cond)] = usrazi[(cond)]
    _usr_comb = (numpy.nansum(usfull * numpy.cos(angle_usr)) * 2 * dazi / 180.)
    _usp_comb = (numpy.nansum(usfull * numpy.sin(angle_usr)) * 2 * dazi / 180.)
    return _usr_comb, _usp_comb


def find_closest(lon: numpy.array, lat: numpy.array, mss: numpy.array,
                 mss_est: numpy.array, ice: numpy.array, hs: numpy.array,
                 beam_angle: list, Rearth: float=6378. * 10**3
                 ) -> Tuple[numpy.ndarray, numpy.ndarray]:
    ''' Find closest mss at 6degree or nadir, closest hs at nadir.
    It assumes that there is a nadir measurement.
    ARGS: 
        lon: longitude for all beams, nadir included
        lat: latitude for all beams, nadir included
        mss: mss measurement for all beams
        mss_est:
        ice: Percentage of ice for all beams
        hs: significant height at nadir
        beam_angle: list anf beam angle
    RETURNS: mssclose: closest mss measurement found
             hsclose: closest hs nadir measurement
    '''
    lon_nadir = lon[:, 0]
    lat_nadir = lat[:, 0]
    lon = lon[:, 1:]
    lat = lat[:, 1:]
    if numpy.min(abs(lon)) <1 and numpy.max(abs(lon))>359:
        lon = numpy.mod(lon + 180, 360) - 180
        lon_nadir = numpy.mod(lon_nadir + 180, 360) - 180
    nsample, nbeam = numpy.shape(mss[:, 1:])
    mssclose = numpy.full((nsample, nbeam), numpy.nan)
    hsclose = numpy.full((nsample, nbeam), numpy.nan)
    ind_6 = numpy.where(numpy.array(beam_angle) == 6)[0]
    mss6 = mss[:, (ind_6 + 1)].ravel()
    mss_nadir = mss[:, 0]
    ice_nadir = ice[:, 0]
    mss_extr = + mss[:, 1:]
    ice_extr = + ice[:, 1:]
    mss_est[numpy.where(ice_extr>0)] = mss_extr[numpy.where(ice_extr>0)]
    ice6 = ice[:, (ind_6 + 1)].ravel()
    lon_nadir_ocean = +lon_nadir
    lon_nadir_ocean[numpy.where((ice_nadir>0))] = numpy.nan # | ~numpy.isfinite(mss_nadir))] = numpy.nan
    lon_nadir_ice = +lon_nadir
    lon_nadir_ice[numpy.where((ice_nadir==0))] = numpy.nan # | ~numpy.isfinite(mss_nadir))] = numpy.nan
    lon_6_ocean = + lon[:, (ind_6)].ravel()
    lon_6_ocean[numpy.where((ice6 > 0))] = numpy.nan #  | ~numpy.isfinite(mss6))] = numpy.nan
    lon_6_ice = + lon[:, (ind_6)].ravel()
    lon_6_ice[numpy.where((ice6 == 0))] = numpy.nan #| ~numpy.isfinite(mss6))] = numpy.nan

    for isample in range(nsample):
        for ibeam in range(nbeam):
            if not numpy.isfinite(mss[isample, ibeam + 1]):
                continue
            plon = lon[isample, ibeam]
            plat = lat[isample, ibeam]
            if ice[isample, ibeam + 1] > 0:
                _lon = lon_nadir_ice
            else:
                _lon = lon_nadir_ocean
            dist_nadir = utils.dist_sphere(plon, _lon, plat, lat_nadir, Rearth)
            #dist_nadir[numpy.where(numpy.isnan(mss_nadir))] = numpy.nan
            # TODO remove ugly try except
            try:
                inadir = numpy.nanargmin(dist_nadir)
            except:
                continue
            mss6[mss6==0] = numpy.nan
            mss_nadir[mss_nadir==0] = numpy.nan

            dnadir = dist_nadir[inadir]
            if ice[isample, ibeam + 1] > 0:
                _lon = lon_6_ice
            else:
                _lon = lon_6_ocean
            dist_6 = utils.dist_sphere(plon, _lon, plat,
                                       lat[:, (ind_6)].ravel(), Rearth)
            try:
                i6 = numpy.nanargmin(dist_6)
            except:
                continue


            d6 = dist_6[i6]
            if d6 > 15 and dnadir > 15:
                mssclose[isample, ibeam] = mss_est[isample, ibeam]
            ind_dist_nadir = numpy.where(dist_nadir < 150)
            ind_dist_6 = numpy.where(dist_6 < 50)
            _hs = numpy.nanmean(hs[ind_dist_nadir])
            # if numpy.isnan(_hs):
            #    continue
            if abs(hs[inadir] - _hs) > 1:
                continue
            #if abs(hs[inadir] - _hs) > 0.4:
            #    hsclose[isample, ibeam] = _hs
            #else:
            hsclose[isample, ibeam] = hs[inadir]
            _mss =  numpy.nanmean(mss6[ind_dist_6])
            if abs(mss6[i6] - _mss) > 0.01:
                continue
            #if abs(mss6[i6] - _mss) > 0.005:
            #    mssclose[isample, ibeam] = _mss
            #else:
            if d6 < dnadir:
                mssclose[isample, ibeam] = mss6[i6]
            else:
                mssclose[isample, ibeam] = mss_nadir[inadir]
            if mss[isample, ibeam + 1] == 0:
                mssclose[isample, ibeam] = 0
    return mssclose, hsclose


def simulate_spect_skim(sgrid: dict, dic_input: dict, beam_angle: list,
                        radial_angle: numpy.array, delta_azim: float=15,
                        fc: int=1, nseed: int=1
                        )-> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
    lon = numpy.transpose(numpy.array(sgrid['lon']))
    lat = numpy.transpose(numpy.array(sgrid['lat']))
    _angle = numpy.array(sgrid['angle'])
    if (sgrid['ipass'] %2) != 0:
       _angle = _angle + numpy.pi
    incl = numpy.transpose(numpy.array(sgrid['incl']))
    list_compulsory = ['uuss', 'vuss', 'mssx', 'mssy', 'mssxy', 'uwnd', 'vwnd',
                       ]
    wnd_dir = numpy.arctan2(dic_input['vwnd'], dic_input['uwnd'])
    wnd_dir = numpy.mod(wnd_dir[:, fc:], 2*numpy.pi)
    usr = utils.proj_radial(dic_input['uuss'][:, fc:],
                                dic_input['vuss'][:, fc:], radial_angle)
    if 'ice' in dic_input.keys():
        ice = dic_input['ice']
    else:
        ice = dic_input['uuss'] * 0
    if 'hs' in dic_input.keys():
        hs = dic_input['hs'][:, 0]
    else:
        hs = dic_input['uuss'][:, 0] * 0
    mss = dic_input['mssu'] + dic_input['mssc']
    _res = combine_usr(lon[:, fc:], lat[:, fc:], usr, dic_input['mssx'][:, fc:],
                       dic_input['mssy'][:, fc:], dic_input['mssxy'][:, fc:],
                       delta_azim, _angle, incl, wnd_dir, nseed=nseed)
    usr_comb, usp_comb, mssr_comb = _res
    mssclose, hsclose = find_closest(lon, lat, mss, mssr_comb, ice, hs,
                                     beam_angle)
    return usr_comb, mssclose, hsclose
