"""
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 sys
import numpy
import pickle
from typing import Dict, Tuple
from scipy import interpolate
from . import utils
import logging
logger = logging.getLogger(__name__)


def compute_rain(time: float, sx_al: numpy.ndarray, sx_al_nadir: numpy.ndarray,
                 sx_ac: numpy.ndarray, dic: dict, nseed: int
                 )-> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray,
                 numpy.ndarray]:
    hour = int((time - numpy.floor(time))*24)
    size_dic = len(dic['xal'][hour])
    numpy.random.seed(nseed)
    rr_ind = int(numpy.random.random_sample() * size_dic)
    xal = dic['xal'][hour][rr_ind]
    var = dic['rr'][hour][rr_ind]
    var2 = dic['pia'][hour][rr_ind]
    xac = dic['xac'][hour][rr_ind]
    x_al_g_tot = + sgrid.x_al
    for i in range(numpy.shape(x_al)[1]):
        x_al_g_tot[:, i] = sx_al[:, i] + sx_al_nadir
    xal_g = numpy.mod(x_al_g_tot - numpy.min(x_al_g_tot), numpy.max(xal))
    interp = interpolate.RectBivariateSpline
    _Teval = interp(xal, xac, numpy.isnan(var), kx=1, ky=1, s=0)
    Teval = _Teval.ev(xal_g, sx_ac)
    # Trick to avoid nan in interpolation
    var_mask = + var
    var_mask[numpy.isnan(var_mask)] = 0.
    var2_mask = + var2
    var2_mask[numpy.isnan(var2_mask)] = 0.
    # Interpolate variable
    _var_out = interp(xal, xac, var_mask, kx=1, ky=1, s=0)
    var_out = _var_out.ev(xal_g, sx_ac)
    _var2_out = interp(xal, xac, var2_mask, kx=1, ky=1, s=0)
    var2_out = _var2_out.ev(xal_g, sx_ac)
    # Mask variable with Teval
    var2_out[Teval > 0] = numpy.nan
    var_out[Teval > 0] = numpy.nan
    xal_n = numpy.mod(sx_al_nadir - numpy.min(sx_al_nadir),
                      numpy.max(xal))
    var_nad = numpy.zeros(numpy.shape(xal_n))
    return var_out, var_nad, var2_out, var_nad


def load_rain(rain_file: str) -> dict:
    with open(rain_file, 'rb') as frain:
        dic = pickle.load(frain)
    return dic


class Rain():
    '''Class that simulate random rain and pia from pyo scenes for all beams
    Args:
        parameters (nseed, list_angle, rain_file, vsat)'''
    def __init__(self, parameters):
        if parameters.rain_file is not None:
            self.rain_dic = load_rain(parameters.rain_file)
        else:
            self.rain_dic = {}
        self.nseed = parameters.nseed + 3
        self.beam_angle = parameters.list_angle
        self.nbeam = len(self.beam_angle)
        self.vsat = parameters.vsat

    def generate(self, time: numpy.ndarray, x_al: numpy.ndarray,
                 x_al_nadir: numpy.ndarray, x_ac: numpy.ndarray,
                 ac_angle: numpy.ndarray) -> Dict[str, numpy.ndarray]:
        ''' Generate rain and gpia for an ensemble of beams as well as nadir
        Args:
            time: time # TODO check format
            x_al: along track distance from nadir
            x_al_nadir: along track distance since beginning of pass
            x_ac: across track distance from nadir
            ac_angle: Angle from the across track direction
        '''

        mean_time = numpy.mean(time)
        rain_all = []
        gpia_err_all = []
        if self.rain_dic:
            rain, rain_nad, gpia, gpia_nad = compute_rain(mean_time,  x_al,
                                                          x_al_nadir, x_ac,
                                                          self.rain_dic,
                                                          self.nseed)
            rain_all.append(rain_nad)
            gpia_err_all.append(gpia_nad)
            for i in range(self.nbeam):
                beam_angle = self.list_angle[i]
                gpia_err = utils.convert_dbkm2ms(gpia, ac_angle[i],
                                                 self.list_angle[i], self.vsat)
                rain_all.append(rain[:, i])
                gpia_err_all.append(gpia_err)
            rain_all = numpy.transpose(numpy.array(rain_all))
            gpia_err_all = numpy.transpose(numpy.array(gpia_err_all))
            return {'rain': rain_all, 'gsig_atm_err': gpia_err_all}
        else:
            return {}


def mask(var: numpy.ndarray, rain: numpy.ndarray, gpia: numpy.ndarray,
         rain_thresh: float) -> numpy.ndarray:
    gpia_thresh = 1
    var[rain > rain_thresh] = numpy.nan
    var[abs(gpia) > gpia_thresh] = numpy.nan
    return var
