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

from sup3r.models.base import Sup3rGan


logger = logging.getLogger(__name__)


class SolarCC(Sup3rGan):
    """Solar climate change model.

    Modifications to standard Sup3rGan:
        - Content loss is only on the n_days of the center 8 daylight hours of
          the daily true+synthetic high res samples
        - Discriminator only sees n_days of the center 8 daylight hours of the
          daily true high res sample.
        - Discriminator sees random n_days of 8-hour samples of the daily
          synthetic high res sample.
    """

    # starting hour is the hour that daylight starts at, daylight hours is the
    # number of daylight hours to sample, so for example if 8 and 8, the
    # daylight slice will be slice(8, 16). The stride length is the step size
    # for sampling the temporal axis of the generated data to send to the
    # discriminator for the adverserial loss component of the generator. For
    # example, if the generator produces 24 timesteps and stride is 4 and the
    # daylight hours is 8, slices of (0, 8) (4, 12), (8, 16), (12, 20), and
    # (16, 24) will be sent to the disc.
    STARTING_HOUR = 8
    DAYLIGHT_HOURS = 8
    STRIDE_LEN = 4

    def init_weights(self, lr_shape, hr_shape, device=None):
        """Initialize the generator and discriminator weights with device
        placement.

        Parameters
        ----------
        lr_shape : tuple
            Shape of one batch of low res input data for sup3r resolution. Note
            that the batch size (axis=0) must be included, but the actual batch
            size doesnt really matter.
        hr_shape : tuple
            Shape of one batch of high res input data for sup3r resolution.
            Note that the batch size (axis=0) must be included, but the actual
            batch size doesnt really matter.
        device : str | None
            Option to place model weights on a device. If None,
            self.default_device will be used.
        """

        # The high resolution data passed to the discriminator should only have
        # daylight hours in the temporal axis=3
        if hr_shape[3] != self.DAYLIGHT_HOURS:
            hr_shape = hr_shape[0:3] + (self.DAYLIGHT_HOURS,) + hr_shape[-1:]

        super().init_weights(lr_shape, hr_shape, device=device)

    @tf.function
    def calc_loss(self, hi_res_true, hi_res_gen, weight_gen_advers=0.001,
                  train_gen=True, train_disc=False):
        """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.
        weight_gen_advers : float
            Weight factor for the adversarial loss component of the generator
            vs. the discriminator.
        train_gen : bool
            True if generator is being trained, then loss=loss_gen
        train_disc : bool
            True if disc is being trained, then loss=loss_disc

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

        if hi_res_gen.shape != hi_res_true.shape:
            msg = ('The tensor shapes of the synthetic output {} and '
                   'true high res {} did not have matching shape! '
                   'Check the spatiotemporal enhancement multipliers in your '
                   'your model config and data handlers.'
                   .format(hi_res_gen.shape, hi_res_true.shape))
            logger.error(msg)
            raise RuntimeError(msg)

        msg = ('Special SolarCC model can only accept multi-day hourly '
               '(multiple of 24) high res data in the axis=3 position but '
               'received shape {}'.format(hi_res_true.shape))
        assert hi_res_true.shape[3] % 24 == 0

        msg = ('Special SolarCC model can only accept multi-day hourly '
               '(multiple of 24) high res synthetic data in the axis=3 '
               'position but received shape {}'.format(hi_res_gen.shape))
        assert hi_res_gen.shape[3] % 24 == 0

        t_len = hi_res_true.shape[3]
        n_days = int(t_len // 24)
        day_slices = [slice(self.STARTING_HOUR + x,
                            self.STARTING_HOUR + x + self.DAYLIGHT_HOURS)
                      for x in range(0, 24 * n_days, 24)]

        # sample only daylight hours for disc training and gen content loss
        disc_out_true = []
        disc_out_gen = []
        loss_gen_content = 0.0
        for tslice in day_slices:
            disc_t = self._tf_discriminate(hi_res_true[:, :, :, tslice, :])
            gen_c = self.calc_loss_gen_content(hi_res_true[:, :, :, tslice, :],
                                               hi_res_gen[:, :, :, tslice, :])
            disc_out_true.append(disc_t)
            loss_gen_content += gen_c

        # Randomly sample daylight windows from generated data. Better than
        # strided samples covering full day because the random samples will
        # provide an evenly balanced training set for the disc
        logits = [[1.0] * (t_len - self.DAYLIGHT_HOURS)]
        time_samples = tf.random.categorical(logits, len(day_slices))
        for i in range(len(day_slices)):
            t0 = time_samples[0, i]
            t1 = t0 + self.DAYLIGHT_HOURS
            disc_g = self._tf_discriminate(hi_res_gen[:, :, :, t0:t1, :])
            disc_out_gen.append(disc_g)

        disc_out_true = tf.concat([disc_out_true], axis=0)
        disc_out_gen = tf.concat([disc_out_gen], axis=0)
        loss_disc = self.calc_loss_disc(disc_out_true, disc_out_gen)

        loss_gen_content /= len(day_slices)
        loss_gen_advers = self.calc_loss_gen_advers(disc_out_gen)
        loss_gen = (loss_gen_content + weight_gen_advers * loss_gen_advers)

        loss = None
        if train_gen:
            loss = loss_gen
        elif train_disc:
            loss = loss_disc

        loss_details = {'loss_gen': loss_gen,
                        'loss_gen_content': loss_gen_content,
                        'loss_gen_advers': loss_gen_advers,
                        'loss_disc': loss_disc,
                        }

        return loss, loss_details
