# -*- coding: utf-8 -*-
"""Sup3r model software"""
import logging
import tensorflow as tf

from sup3r.models.base import Sup3rGan


logger = logging.getLogger(__name__)


class WindCC(Sup3rGan):
    """Wind climate change model.

    Modifications to standard Sup3rGan:
        - Hi res topography is expected as the last feature channel in the true
          data in the true batch observation. This topo channel is appended to
          the generated output so the discriminator can look at the wind fields
          compared to the associated hi res topo.
    """

    def set_model_params(self, **kwargs):
        """Set parameters used for training the model

        Parameters
        ----------
        kwargs : dict
            Keyword arguments including 'training_features', 'output_features',
            'smoothed_features', 's_enhance', 't_enhance', 'smoothing'
        """
        output_features = kwargs['output_features']
        msg = ('Last output feature from the data handler must be topography '
               'to train the WindCC model, but received output features: {}'
               .format(output_features))
        assert output_features[-1] == 'topography', msg
        output_features.remove('topography')
        kwargs['output_features'] = output_features
        super().set_model_params(**kwargs)

    @tf.function
    def calc_loss(self, hi_res_true, hi_res_gen, **kwargs):
        """Calculate the GAN loss function using generated and true high
        resolution data.

        Parameters
        ----------
        hi_res_true : tf.Tensor
            Ground truth high resolution spatiotemporal data.
        hi_res_gen : tf.Tensor
            Superresolved high resolution spatiotemporal data generated by the
            generative model.
        kwargs : dict
            Key word arguments for:
            Sup3rGan.calc_loss(hi_res_true, hi_res_gen, **kwargs)

        Returns
        -------
        loss : tf.Tensor
            0D tensor representing the loss value for the network being trained
            (either generator or one of the discriminators)
        loss_details : dict
            Namespace of the breakdown of loss components
        """

        # append the true topography to the generated synthetic wind data
        hi_res_gen = tf.concat((hi_res_gen, hi_res_true[..., -1:]), axis=-1)

        return super().calc_loss(hi_res_true, hi_res_gen, **kwargs)
