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

"""
Generate instrumental errors
----------------------------
"""
from typing import Dict
import numpy
import datetime
from . import (Altimeter, WaveDoppler, Instrument, Rain, Attitude, DSigma)
from . import utils
# WetTroposphere


class Generator:
    """Instrumental error generator.

    Args:
        parameters (settings.Parameters): Simulation settings
        first_date (numpy.datetime64): Date of the first simulated
            measurement.
    """
    def __init__(self, parameters):
        #: The list of user-defined error generators
        self.generators = []

        for item in parameters.noise:
            if item == Altimeter.__name__:
                self.generators.append(Altimeter(parameters))
            elif item == Attitude.__name__:
                self.generators.append(Attitude(parameters))
            elif item == DSigma.__name__:
                self.generators.append(DSigma(parameters))
            elif item == Instrument.__name__:
                self.generators.append(Instrument(parameters))
            elif item == Rain.__name__:
                self.generators.append(Rain(parameters))
            elif item == WaveDoppler.__name__:
                self.generators.append(WaveDoppler(parameters))
            # elif item == WetTroposphere.__name__:
            #    self.generators.append(WetTroposphere(parameters))
            else:
                # A new error generation class has been implemented but it is
                # not handled by this object.
                raise ValueError(f"unknown error generation class: {item}")

    def generate_dask(self, cycle_number: int, curvilinear_distance: float,
                 time: numpy.array, x_al: numpy.array, x_ac: numpy.array,
                 swh: numpy.array) -> Dict[str, numpy.ndarray]:
        """Generate errors

        Args:
            cycle_number (int): Cycle number.
            curvilinear_distance (float): Curvilinear distance covered by the
                satellite during a complete cycle.
            time (numpy.ndarray): Date of measurements.
            x_al (numpy.ndarray): Along track distance.
            x_ac (numpy.ndarray): Across track distance.

        Returns:
            dict: Associative array between error variables and simulated
            values.
        """
        result = {}
        if not self.generators or x_al.shape[0] == 0:
            return result

        futures = []
        with dask.distributed.worker_client() as client:
            for item in self.generators:
                if isinstance(item, Altimeter):
                    futures.append(client.submit(item.generate, x_al))
                elif isinstance(item, BaselineDilation):
                    futures.append(client.submit(item.generate, x_al, x_ac))
                elif isinstance(item, CorrectedRollPhase):
                    futures.append(client.submit(item.generate, time, x_ac))
                elif isinstance(item, Karin):
                    futures.append(
                        client.submit(item.generate, x_al, x_ac,
                                      curvilinear_distance, cycle_number, swh))
                elif isinstance(item, RollPhase):
                    futures.append(client.submit(item.generate, x_al, x_ac))
                elif isinstance(item, Timing):
                    futures.append(client.submit(item.generate, x_al, x_ac))
                elif isinstance(item, WetTroposphere):
                    futures.append(client.submit(item.generate, x_al, x_ac))

            for future in dask.distributed.as_completed(futures):
                result.update(future.result())
        return result


    def generate(self, x_al: numpy.array, dic_input: dict,
                 ac_angle: numpy.array, time:numpy.array, ti_cycle: float,
                 first_time: datetime.datetime, radial_angle: numpy.array,
                 x_al_nadir: numpy.array, x_ac: numpy.array
                 ) -> Dict[str, numpy.ndarray]:
        res = {}
        fc = 1
        for item in self.generators:
            if isinstance(item, Altimeter):
                res.update(item.generate(x_al[:, 0]))
            elif isinstance(item, DSigma):
                sigma0 = dic_input['sigma0']
                res.update(item.generate(sigma0, ac_angle))
            elif isinstance(item, Attitude):
                res.update(item.generate(time, ti_cycle, ac_angle, first_time))
            elif isinstance(item, Instrument):
                sigma0 = dic_input['sigma0']
                wind_dir = numpy.arctan2(dic_input['vwnd'][:, fc:],
                                         dic_input['uwnd'][:, fc:])
                wind_norm = numpy.sqrt(dic_input['uwnd'][:, fc:]**2
                                       + dic_input['vwnd'][:, fc:]**2)
                res.update(item.generate(radial_angle, ac_angle, wind_norm,
                                         wind_dir, sigma0))
            elif isinstance(item, Rain):
                res.update(item.generate(time, x_al, x_al_nadir, x_ac,
                                         ac_angle))
            elif isinstance(item, WaveDoppler):
                sigma0 = dic_input['sigma0']
                wind_norm = numpy.sqrt(dic_input['uwnd'][:, fc:]**2
                                       + dic_input['vwnd'][:, fc:]**2)
                windr = utils.proj_radial(dic_input['uwnd'][:, fc:],
                                              dic_input['vwnd'][:, fc:],
                                              radial_angle)
                wind_dir = numpy.arctan2(dic_input['vwnd'][:, fc:],
                                         dic_input['uwnd'][:, fc:])
                if 'mssu' in dic_input.keys():
                    mss = dic_input['mssu'][:, fc:] + dic_input['mssc'][:, fc:]
                else:
                    logger.error('mss variable not found')
                    sys.exit(1)
                if 'hs' in dic_input.keys():
                    hs = dic_input['hs'][:, fc:]
                else:
                    hs = sigma0 * 0

                if 'uuss' in dic_input.keys():
                    uuss = dic_input['uuss'][:, fc:]
                    vuss = dic_input['vuss'][:, fc:]
                else:
                    uuss = sigma0 * 0
                    vuss = sigma0 * 0
                if 'ice' in dic_input.keys():
                    ice = dic_input['ice'][:, fc:]
                else:
                    ice = sigma0 * 0
                res.update(item.generate(wind_norm, wind_dir, windr, mss, hs,
                                         ice, sigma0, radial_angle, uuss, vuss))
        return res


    def correct(self, grid_dict: dict, dic_input: dict,
                radial_angle: numpy.array, dazim: float, fc: int=1
                ) -> Dict[str, numpy.ndarray]:
        res = {}
        for item in self.generators:
            if isinstance(item, WaveDoppler):
                sigma0 = + dic_input['sigma0']
                wind_norm = numpy.sqrt(dic_input['uwnd'][:, fc:]**2
                                       + dic_input['vwnd'][:, fc:]**2)
                windr = utils.proj_radial(dic_input['uwnd'][:, fc:],
                                              dic_input['vwnd'][:, fc:],
                                              radial_angle)
                if 'ice' in dic_input.keys():
                    ice = dic_input['ice'][:, fc:]
                else:
                    ice = sigma0 * 0
                res.update(item.correct(grid_dict, dic_input, wind_norm, windr,
                                        ice, sigma0, radial_angle, dazim))
        return res

    def generate_obs(self, dic_input: dict, error:dict, corr: dict, fc: int=1,
                    ) -> Dict[str, numpy.ndarray]:
        ur_obs = + dic_input['ur_true']
        nsample = numpy.shape(dic_input['ssh_true'])[0]
        ssh_obs = + dic_input['ssh_true'][:, 0].reshape(nsample, -1)
        for item in self.generators:
            if isinstance(item, Altimeter):
                fc = 1
                ssh_obs += error['altimeter']
            elif isinstance(item, DSigma):
                ur_obs += error['dsigma']
            elif isinstance(item, Attitude):
                # TODO add attitude corrected error or yaw total
                att = 1
            elif isinstance(item, Instrument):
                ur_obs += error['instr']
            elif isinstance(item, Rain):
                # TODO mask data
                mask = fc
            elif isinstance(item, WaveDoppler):
                if 'wd_est' in dic_input.keys():
                    wd_corr = (error['wd'] - corr['wd_est'])
                else:
                    wd_corr = error['wd']
                ur_obs += wd_corr
        return {"ur_obs": ur_obs, "ssh_obs": ssh_obs}
