"""Tensorflow implementation of adversarial autoencoders.

Originally based on [1] with significant changers.

Demo (command line)
===================

conda activate tf114
python -m dltb.thirdparty.tensorflow.aae

Demo (interactive)
==================

# import tensorflow as tf
import tensorflow.compat.v1 as tf
from ulf import SupervisedAdversarialAutoencoder
aae = SupervisedAdversarialAutoencoder(code_dim=2)
aae.prepare()

saver = tf.train.Saver(name='aae', filename="my_test_model")
sess = tf.InteractiveSession()
aae.set_tensorflow_session(sess)
start_epoch = aae._tf_initialize_variables(sess, saver)

aae.plot_data_codes_2d(test_xs, labels=test_ys)

aae.plot_recoded_images(test_xs[:100], labels=test_ys[:100])

aae.plot_decoded_codespace_2d(labels=3)

aae.plot_analogical_decoding()

sess.close()

Data
====
# data are numpy.ndarray.
# xs: image pixel data
#     dtype=float64, min/max=0.0/1.0,
#     shape=(batch, height, width, channels)
# ys: one-hot encoded class labels
#     dtype=float64, min/max=0.0/1.0
#     shape=(batch, classes)

      sellf._data_semi_pipeline = data_pipeline(self._conf.data)
      self._valid_xs = self._valid_xs[:self._conf.num_samples]
      self._valid_ys = self._valid_ys[:self._conf.num_samples]

References
==========
[1] https://github.com/MINGUKKANG/Adversarial-AutoEncoder.git
"""

# pylint: disable=too-many-lines
# pylint: disable=fixme
# pylint: disable=unexpected-keyword-arg,no-value-for-parameter

# standard imports
from typing import Tuple, Optional, Sized, Union
import time
import datetime

# third-party imports
import numpy as np
from tqdm import tqdm

# toolbox imports
from dltb.datasource import Datasource
from dltb.base.data import add_noise
from dltb.util.plot import TilingPlotter
from dltb.util.distributions import gaussian, gaussian_mixture, swiss_roll
from . import v1 as tf
from .v1 import Utils as tf_helper
from .ae import Autoencoder

class AdversarialAutoencoder(Autoencoder):
    """Base class for the TensorFlow Adversarial Autoencoder (AAE)
    implementation.  The AAE extends the basic Autoencoder (AE) by
    introducing an additional training objective to force the encoder
    to produces code that is distributed according to a given target
    distribution ("prior").

    The method :py:class:`_sample_prior` allows to samples codes
    from the target distribution.  There are three different target
    distributions implemented by this class: `'gaussian'`,
    `'gaussian_mixture'`, and `'swiss_roll'`.

    Training
    --------
    Training the AAE consists of two parts: on the one hand, the full
    autoencoder stack (encoder + decoder) is trained to minimize the
    reconstruction error.  On the other hand, the encoder part is
    trained to fit the code distribution to a given target distribution
    ("prior").  This training employs an adversarial training procedure
    in which the encoder acts as generator that tries to mimic the
    target distribution.

    The autoencoder training aims to minimize the reconstruction loss,
    which is accessible via the property `_tf_loss_reconstruction`.  It is
    interpreted as the "negative log likelihood".  Computing this loss
    only requires the input data values.

    The adversarial training process consist of two parts: on the one
    hand it aims to improve the discriminator in its ability to
    discriminate codes output by the encode from codes sampled from
    the target distribution. On the other hand it wants the encoder to
    improve to create codes that follow the target distribution.

    Arguments
    ---------
    data:
        The dataset to use (either 'MNIST' or '')

    """

    def __init__(self, data: str = 'MNIST',
                 prior: str = 'gaussian', **kwargs) -> None:
        super().__init__(**kwargs)
        self._conf.data = data
        self._conf.prior = prior

        # tensorflow placeholders
        self._tf_prior = None

        # The model
        self._tf_loss_discriminator = None
        self._tf_loss_generator = None

        # Optimizers
        self._tf_optimize_discriminator = None
        self._tf_optimize_generator = None

        # loss values (per batch)
        self._loss_discriminator_value = None
        self._loss_generator_value = None

    def _prepare_tensors(self) -> None:
        """
        """
        if self._tf_prior is None:
            self._tf_prior = \
                tf.placeholder(tf.float32, shape=[None, self.code_dim],
                               name="z_prior")

        super()._prepare_tensors()

    @staticmethod
    def _sample_prior(prior: str, zdim: int, nclasses: int, batch_size: int,
                      use_label_info: bool) -> Tuple[np.ndarray, np.ndarray]:
        """
        Results
        -------
        prior_style:
        prior_label_onehot:
            A 2-dimensional numpy.ndarray of type float holding `prior_label`
            in one-hot-encoding. The shape is (batch_size, nclasses).
        """
        if prior == 'gaussian':
            prior_style, prior_label = \
                gaussian(batch_size, labels=nclasses)

        elif prior == 'gaussian_mixture':
            prior_style, prior_label = \
                gaussian_mixture(batch_size, components=nclasses, labels=True)

        elif prior == 'swiss_roll':
            prior_style, prior_label = \
                swiss_roll(batch_size, labels=nclasses)

        else:
            raise ValueError(f"Unknow prior: '{prior}'. Known values are: "
                             "gaussian, gaussian_mixture, swiss_roll")

        prior_label_onehot = np.eye(nclasses, dtype=np.float32)[prior_label]
        return prior_style, prior_label_onehot

    def _tf_define_optimizers(self, variables) -> None:
        """Prepare the training process. Define optimizers.
        """
        super()._tf_define_optimizers(variables)

        var_generator = [var for var in variables
                         if "encoder" in var.name]
        var_discriminator = [var for var in variables
                             if "discriminator" in var.name]

        self._tf_optimize_discriminator = \
            tf.train.AdamOptimizer(learning_rate=self._tf_learning_rate/5).\
            minimize(self._tf_loss_discriminator,
                     global_step=self._tf_global_step,
                     var_list=var_discriminator)
        self._tf_optimize_generator = \
            tf.train.AdamOptimizer(learning_rate=self._tf_learning_rate).\
            minimize(self._tf_loss_generator,
                     global_step=self._tf_global_step,
                     var_list=var_generator)

    def _tf_train_step(self, sess, feed_dict) -> None:
        """Perform one training step.
        """
        super()._tf_train_step(sess, feed_dict)
        # Optmize the reconstruction loss. This requires
        # 'inputs' and 'outputs' to be given in the 'feed_dict'
        # _negative_log_likelihood_value, _, _g = \
        self._loss_reconstruction_value, _, _g = \
            sess.run([self._tf_loss_reconstruction,
                      self._tf_optimize_reconstruction,
                      self._tf_global_step],
                     feed_dict=feed_dict)

        # Discriminator phase
        self._loss_discriminator_value, _ = \
            sess.run([self._tf_loss_discriminator,
                      self._tf_optimize_discriminator],
                     feed_dict=feed_dict)

        # Generator phase
        self._loss_generator_value, _ = \
            sess.run([self._tf_loss_generator,
                      self._tf_optimize_generator],
                     feed_dict=feed_dict)

class LabeledAutoencoder(Autoencoder):
    """A labeled autoencoder splits the code interpretation into two
    parts: one part representing the class label (as one-hot encoded
    vector) while the other part should hold other information
    (sometimes referred to as "style").

    A labeled autoencoder provides methods to compute and access the
    two code parts (style and label) independently.

    In a `LabeledAutoencoder`, the property :py:prop:`code_dim` refers
    to the dimensionality of the style part of the code, while the
    property :py:prop:`nclasses` contains the number of classes.
    That is, the combined code vector (style + label) has dimensionality
    `code_dim + nclasses`.
    """

    def __init__(self, nclasses: Optional[int] = None, **kwargs) -> None:
        super().__init__(**kwargs)

        # properties
        self._nclasses = nclasses

        # tensorflow placeholders
        self._tf_labels = None

    @property
    def nclasses(self) -> int:
        """The number of classes, that is the number of different labels,
        in this :py:class:`LabeledAutoencoder`.
        """
        return self._nclasses

    def _prepare_tensors(self) -> None:
        """
        """
        self._tf_labels = \
            tf.placeholder(dtype=tf.float32, shape=[None, self.nclasses],
                           name="Input_labels")

        super()._prepare_tensors()

    def one_hot(self, labels: np.ndarray,
                length: Optional[int] = None) -> np.ndarray:
        """Get the given labels in one-hot encoding.
        """
        if isinstance(labels, Sized):
            if length is not None and len(labels) != length:
                raise ValueError(f"Labels have length {len(labels)} but "
                                 f"should have lenght {length}.")
            if isinstance(labels, np.ndarray):
                if labels.ndim == 2:
                    return labels
                if labels.ndim == 1:
                    return np.eye(self.nclasses, dtype=np.float32)[labels]
            elif isinstance(labels, list):
                return np.eye(self.nclasses, dtype=np.float32)[labels]
        elif isinstance(labels, int):
            if length is None:
                raise ValueError("You need to specify a length for "
                                 "the one-hot vector")
            return np.eye(self.nclasses, dtype=np.float32)[[labels] * length]
        raise TypeError("Unexpected type {type(labels)} for labels.")

class LabeledAdversarialAutoencoder(LabeledAutoencoder,
                                    AdversarialAutoencoder):
    """A labeled adversarial autoencoder combines the ideas of the labeled
    autoencoder and the adversarial autoencoder.  It aims at forcing
    both, the distribution of the style part of code and the
    distribution of the label part of the code towards two given
    priors using adversarial training techniques.
    """

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

        # tensorflow placeholders
        self._tf_prior_style = None
        self._tf_prior_label = None

    def _prepare_tensors(self) -> None:
        """
        """
        self._tf_prior_style = \
            tf.placeholder(tf.float32, shape=[None, self.code_dim],
                           name="z_prior")
        self._tf_prior_label = \
            tf.placeholder(tf.float32, shape=[None, self.nclasses],
                           name="prior_labels")

        self._tf_prior = \
            tf.concat([self._tf_prior_style, self._tf_prior_label], axis=1)

        super()._prepare_tensors()

    def plot_analogical_decoding(self, style: Union[np.ndarray, int] = 10,
                                 plotter: Optional[TilingPlotter] = None,
                                 **kwargs) -> None:
        """Plot the analogical reasoning results. For a given style,
        generate analogical exemplas for each of the classes.

        Arguments
        ---------
        style:
            The styles to be used for generation. If an `int`
            this will specify the number of random styles to
            to be choosen.
        """
        # generate random code style samples (styles)
        if isinstance(style, int):  # number of examples
            style = np.random.rand(style, self.code_dim)
        examples = len(style)
        style = np.repeat(style, self.nclasses, axis=0)

        labels = np.arange(self.nclasses)
        labels = np.tile(labels, examples)
        one_hot = self.one_hot(labels)

        codes = np.concatenate((style, one_hot), axis=1)

        # use decoder to generate output data from the combined code.
        data = self.decode(codes)

        # plot the results
        if plotter is None:
            plotter = TilingPlotter()
        plotter.plot_tiling(data, rows=examples, columns=self.nclasses,
                            **kwargs)

class SupervisedAdversarialAutoencoder(LabeledAdversarialAutoencoder):
    """Specialized sublasse for a supervised Adversarial Autoencoder (AAE)
    implementation.

    Notice that the :py:class:`SupervisedAdversarialAutoencoder` is
    not really an autoencoder, as the code does not contain sufficient
    information to reconstruct the input.  Instead the decoder needs
    additional label information.

    The encoder encodes a given input just in the "style" part.  This
    style code has to be combined with the label information to obtain
    the full code.

    The encoder uses `tf_inputs` to compute the style code which is
    provided in `tf_encoded_style`.  This can be combined with
    information from `tf_labels` to obtain the full code vector in
    `tf_encoded`. The decoder uses `tf_encoded` to reconstruct the
    data and provides it in `tf_decoded`.

    This specific behavior results in a slight redesign of the
    autoencoder interface: the methods :py:meth:`encode`,
    :py:meth:`decode`, and :py:meth:`recode` accept an additional
    argument `labels` that allows to provide the required class
    labels.

    """

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self._conf.model = 'supervised'

        self._tf_encoded_style = None

    def encode(self, data: np.ndarray, batch_size: int = 128,
               labels: Optional[np.ndarray] = None) -> np.ndarray:
        """Encode data using the encoder part of the autoencoder.

        Arguments
        ---------
        data:
            The data to be encoded.

        Result
        ------
        code:
            The codes obtained from the data. If no label information
            is available, this will only be the style part of the code.
        """
        length = len(data)
        code = np.ndarray((length, self.code_dim))

        offset = 0
        batches = self._np_batcher(data, batch_size=batch_size)
        code_batches = \
            self._tf_encode_batches(batches, tf_encoded=self._tf_encoded_style)
        for batch in code_batches:
            end = offset + len(batch)
            code[offset: end] = batch
            offset = end

        if labels is None:
            return code

        labels_one_hot = self.one_hot(labels, length=length)
        return np.concatenate((code, labels_one_hot), axis=1)

    def decode(self, code: np.ndarray, batch_size: int = 128,
               labels: Optional[np.ndarray] = None) -> np.ndarray:
        """Decode given code values into the data space using the decoder
        part of the autoencoder.

        Arguments
        ---------
        code:
            The codes to be decoded.

        Result
        ------
        data:
            The reconstructed data.
        """
        if code.shape[1] == self.code_dim:
            if labels is None:
                raise ValueError("The supervised autoencoder needs explicit "
                                 "label information for decoding.")
            labels_one_hot = self.one_hot(labels, length=len(code))
            code = np.concatenate((code, labels_one_hot), axis=1)
        return super().decode(code, batch_size=batch_size)

    def recode(self, data: np.ndarray, batch_size: int = 128,
               labels: Optional[np.ndarray] = None) -> np.ndarray:
        """Reconstruct data values using the autoencoder, that is first
        encode the data and the decode it back into the data space.

        Arguments
        ---------
        data:
            The data to be recoded.

        Result
        ------
        recoded:
            The reconstructed data.
        """
        # FIXME[hack]: just encode + decode. A full tensorflow
        # recoding would probably be more efficient
        code = self.encode(data, batch_size=batch_size, labels=labels)
        return self.decode(code, batch_size=batch_size)

    def _prepare_tensors(self) -> None:
        """Setup TensorFlow properties for the supervised adversarial
        autoencoder.

        Prepare the tensorflow network (computational graph) realizing
        the supervised adversarial autoencoder.

        The constructed graph looks as follows:

                 [inputs]           [labels]
                    |                  |
               inputs_flat             |
                    |                  |
          (encoder) |                  |
                    V                  |
             *encoded_style*           |      [prior_style]  [prior_label]
                    |                  |             |            |
                    +------------------+             +------+-----+
                             |                              |
                             V                              V
                         *encoded*                        prior
                             |                              |
                    +----------------+                      |
          (decoder) |                |                      |
                    |                V                      V
                    V             D_fake_logits        D_real_logits
                *decoded*          |         |              |
                    |              V         V              V
                    +-----+     G_loss   D_loss_fake   D_loss_true
                          |        |           |              |
            [outputs]     |        V           +-------+------+
                |         |   *loss_generator*         V
            outputs_flat  |                   *loss_discriminator*
                V         V
             *loss_reconstruction*

        [...]: input values (tf.placeholder)
                 inputs
                 outputs
                 code
                 prior_style
                 prior_label
        *...*: output values (tf.tensor)
                 encoded: encoder output (just style, no label)
                 decoded: decoder output (flat data)

        Tensorflow properties are prefixed with '_tf_'.
        """
        # Super class sets up a standard autoencoder from placeholders
        # 'inputs', 'outputs', providing 'encoded', 'decoded' and
        # 'loss_reconstruction'.
        super()._prepare_tensors()

        #
        # The discriminators
        #
        prior = \
            tf.concat([self._tf_prior_style, self._tf_prior_label], axis=1)
        discriminator_real_logits = \
            self._tf_discriminator(prior, self._tf_keep_prob)
        discriminator_fake_logits = \
            self._tf_discriminator(self._tf_encoded, self._tf_keep_prob)

        discriminator_fake_labels = tf.zeros_like(discriminator_fake_logits)
        discriminator_loss_fake = tf.nn.\
            sigmoid_cross_entropy_with_logits(logits=discriminator_fake_logits,
                                              labels=discriminator_fake_labels)
        discriminator_real_labels = tf.ones_like(discriminator_real_logits)
        discriminator_loss_true = tf.nn.\
            sigmoid_cross_entropy_with_logits(logits=discriminator_real_logits,
                                              labels=discriminator_real_labels)

        generator_fake_labels = tf.ones_like(discriminator_fake_logits)
        generator_loss = tf.nn.\
            sigmoid_cross_entropy_with_logits(logits=discriminator_fake_logits,
                                              labels=generator_fake_labels)

        self._tf_loss_discriminator = \
            tf.reduce_mean(discriminator_loss_fake) + \
            tf.reduce_mean(discriminator_loss_true)
        self._tf_loss_generator = tf.reduce_mean(generator_loss)

    def _tf_encoder(self, data, keep_prob):
        """Encoder for an autoencoder.
        """
        self._tf_encoded_style = super()._tf_encoder(data, keep_prob)
        return tf.concat([self._tf_encoded_style, self._tf_labels], axis=1)

    def train(self,
              labeled_train_data: Datasource,
              labeled_display_data: Datasource,
              sess: tf.Session, saver) -> None:
        """Train the network.
        """
        # display_data: (images, images_noised, labels)
        # A batch of data used for plotting intermediate results
        # during training (taken from the validation dataset)
        # Each is a numpy array of length 100 and appropriate shape.
        display_data = labeled_display_data[:100, ('array', 'array', 'label')]

        #
        # prepare the optimizers
        #

        # obtain the variables used in the model
        total_vars = tf.trainable_variables()
        # var_ae = [var for var in total_vars
        #           if "encoder" in var.name or "decoder" in var.name]

        # Optimizers
        self._tf_define_optimizers(total_vars)

        #
        # Start the training
        #
        start_epoch = self._tf_initialize_variables(sess, saver)
        start_time = time.time()

        total_batch = len(labeled_train_data) // self._conf.batch_size
        labeled_batch_iterator = \
            labeled_train_data(batch_size=self._conf.batch_size, loop=True,
                               attributes=('array', 'noisy', 'label'))

        for epoch in tqdm(range(start_epoch, self._conf.n_epoch),
                          initial=start_epoch, total=self._conf.n_epoch):
            likelihood = 0
            discriminator_value = 0
            generator_value = 0

            # Adapt the learning rate depending on the epoch
            lr_value = self._learning_rate_schedule(epoch)

            for _batch_idx in tqdm(range(total_batch)):
                batch_xs, batch_noised_xs, batch_ys = \
                    next(labeled_batch_iterator)

                # Sample from the prior distribution
                prior_style, prior_label_onehot = \
                    self._sample_prior(self._conf.prior, zdim=self.code_dim,
                                       nclasses=self.nclasses,
                                       batch_size=self._conf.batch_size,
                                       use_label_info=True)

                feed_dict = {
                    self._tf_inputs: batch_noised_xs,
                    self._tf_outputs: batch_xs,
                    self._tf_labels: batch_ys,
                    self._tf_prior_style: prior_style,
                    self._tf_prior_label: prior_label_onehot,
                    self._tf_learning_rate: lr_value,
                    self._tf_keep_prob: self._conf.keep_prob
                }

                # AutoEncoder phase
                self._tf_train_step(sess, feed_dict)

                # Summary
                likelihood += \
                    self._loss_reconstruction_value/total_batch
                discriminator_value += \
                    self._loss_discriminator_value/total_batch
                generator_value += \
                    self._loss_generator_value/total_batch

            # every 5th epoch (except the last) plot the manifold canvas
            if epoch % 5 == 0 or epoch == (self._conf.n_epoch - 1):
                name = f"Manifold_canvas_{epoch}"
                self.plot_recoded_images(display_data[1],
                                         targets=display_data[0],
                                         labels=display_data[2],
                                         filename=name)

            # output end of epoch information
            runtime = time.time() - start_time
            print(f"Epoch: {epoch:3d}, "
                  f"global step: {sess.run(self._tf_global_step)}, "
                  f"Time: {datetime.timedelta(seconds=runtime)}")
            print(f"             lr_AE: {lr_value:.5f}"
                  f"   loss_AE: {likelihood:.4f}   ")
            print(f"             lr_D: {lr_value/5:.5f}"
                  f"   loss_D: {discriminator_value:.4f}")
            print(f"             lr_G: {lr_value:.5f}"
                  f"   loss_G: {generator_value:.4f}\n")

            if saver is not None:
                saver.save(sess, 'checkpoints/my_test_model',
                           global_step=self._tf_global_step,
                           write_meta_graph=False)
                print(f"Saver: {saver.last_checkpoints}")

class SemisupervisedAdversarialAutoencoder(LabeledAdversarialAutoencoder):
    """Specialized sublasse for a semi-supervised
    Adversarial Autoencoder (AAE) implementation.

    The main differences to the fully supervised AAE are
    the following:
    * the encoder (generator) also outputs class labels. That is the latent
      representation is split into two parts: the continuous z and the one-hot
      encoded label information y.
    * there now are two discriminators, one for each part of the latent
      representation, and two loss functions training them:
      `_loss_discriminator_style` and `_loss_discriminator_label`.
      Training the z part requires z value output from the encoder/generator
      ()
      as well as real z values sampled from the target distribution.
      Training the y part requires the label output from the encoder/generator

    * the produced class labels can be used as additional training objective
      to train the encoder (generator) to minimize crossentropy loss, if
      ground truth label are available (supervised case).  This loss function
      is stored under the name `_crossentropy_labels`.  The training process
      requires input data with real class labels.
    """

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self._conf.model = 'semi_supervised'

        #
        # data related properties
        #

        # model related properties
        self._style = None
        self._crossentropy_labels = None

        # loss functions
        self._tf_loss_discriminator_label = None
        self._tf_loss_discriminator_style = None

        # optimizers
        self._op_z_discriminator = None
        self._op_y_discriminator = None
        self._op_generator = None
        self._op_crossentropy_labels = None

        self._l_z_discriminator = None
        self._l_y_discriminator = None
        self._l_generator = None
        self._crossentropy = None

    def _prepare_tensors(self) -> None:
        """Setup TensorFlow properties for the supervised adversarial
        autoencoder.

        """
        super()._prepare_tensors()

        # placeholders
        # FIXME[semi]: self._z_cat
        # Y_cat = tf.placeholder(dtype=tf.float32,
        #                        shape=[None, n_cls], name="labels_cat")
        # labels_cat = self._tf_prior_label

        # FIXME[coding]: code duplication
        flat_data_length = \
            self._data_shape[0] * self._data_shape[1] * self._data_shape[2]
        inputs_flat = tf.reshape(self._tf_inputs, [-1, flat_data_length])
        outputs_flat = tf.reshape(self._tf_outputs, [-1, flat_data_length])

        # the encoder
        self._style, labels_softmax = \
            self._tf_semi_encoder(inputs_flat, self._tf_keep_prob,
                                  semi_supervised=False)
        _, labels_generated = \
            self._tf_semi_encoder(inputs_flat, self._tf_keep_prob,
                                  semi_supervised=True)
        latent_inputs = tf.concat([self._style, labels_softmax], axis=1)

        # the decoder
        self._tf_decoded = \
            self._tf_semi_decoder(latent_inputs, self._tf_keep_prob)

        #
        # the discriminators
        #

        discriminator_label_fake = \
            self._tf_semi_y_discriminator(labels_softmax,
                                          self._tf_keep_prob)
        discriminator_label_real = \
            self._tf_semi_y_discriminator(self._tf_prior_label,
                                          self._tf_keep_prob)

        discriminator_style_fake = \
            self._tf_semi_z_discriminator(self._style,
                                          self._tf_keep_prob)
        discriminator_style_real = \
            self._tf_semi_z_discriminator(self._tf_prior_style,
                                          self._tf_keep_prob)

        #
        # loss functions
        #
        self._tf_loss_reconstruction = \
            tf.reduce_mean(tf.squared_difference(self._tf_decoded,
                                                 outputs_flat))

        discriminator_label_zeros = tf.zeros_like(discriminator_label_fake)
        discriminator_label_ones = tf.ones_like(discriminator_label_real)
        discriminator_loss_label_real = tf.nn.\
            sigmoid_cross_entropy_with_logits(logits=discriminator_label_real,
                                              labels=discriminator_label_ones)
        discriminator_loss_label_fake = tf.nn.\
            sigmoid_cross_entropy_with_logits(logits=discriminator_label_fake,
                                              labels=discriminator_label_zeros)
        self._tf_loss_discriminator_label = \
            tf.reduce_mean(discriminator_loss_label_real) + \
            tf.reduce_mean(discriminator_loss_label_fake)

        discriminator_style_zeros = tf.zeros_like(discriminator_style_fake)
        discriminator_style_ones = tf.ones_like(discriminator_style_real)
        discriminator_loss_style_real = tf.nn.\
            sigmoid_cross_entropy_with_logits(logits=discriminator_style_real,
                                              labels=discriminator_style_ones)
        discriminator_loss_style_fake = tf.nn.\
            sigmoid_cross_entropy_with_logits(logits=discriminator_style_fake,
                                              labels=discriminator_style_zeros)
        self._tf_loss_discriminator_style = \
            tf.reduce_mean(discriminator_loss_style_real) + \
            tf.reduce_mean(discriminator_loss_style_fake)

        loss_generator_label = tf.nn.\
            sigmoid_cross_entropy_with_logits(logits=discriminator_label_fake,
                                              labels=discriminator_label_ones)
        loss_generator_style = tf.nn.\
            sigmoid_cross_entropy_with_logits(logits=discriminator_style_fake,
                                              labels=discriminator_style_ones)
        self._tf_loss_generator = \
            tf.reduce_mean(loss_generator_style) + \
            tf.reduce_mean(loss_generator_label)

        crossentropy_labels = tf.nn.\
            softmax_cross_entropy_with_logits(logits=labels_generated,
                                              labels=self._tf_labels)
        self._crossentropy_labels = tf.reduce_mean(crossentropy_labels)

    def _tf_semi_encoder(self, data, keep_prob, semi_supervised=False):
        """Encoder for semi-supervised AAE.

        Arguments
        ---------
        """
        with tf.variable_scope("semi_encoder", reuse=tf.AUTO_REUSE):
            net = tf_helper.dense_layer(data, self._conf.semi_n_hidden,
                                        name="dense_1", keep_prob=keep_prob)
            net = tf_helper.dense_layer(net, self._conf.semi_n_hidden,
                                        name="dense_2", keep_prob=keep_prob)
            style = tf_helper.dense(net, self.code_dim, name="style")

            if semi_supervised is False:
                labels_generated = \
                    tf.nn.softmax(tf_helper.dense(net, self.nclasses,
                                                  name="labels"))
            else:
                labels_generated = \
                    tf_helper.dense(net, self.nclasses,
                                    name="label_logits")

        return style, labels_generated

    def _tf_semi_decoder(self, code, keep_prob):
        """Decoder for semi-supervised AAE.

        Result
        ------
        decoder:
            A flat tensor holding the decoded data.
        """
        flat_data_length = \
            self._data_shape[0] * self._data_shape[1] * self._data_shape[2]
        with tf.variable_scope("semi_decoder", reuse=tf.AUTO_REUSE):
            net = tf_helper.dense_layer(code, self._conf.semi_n_hidden,
                                        name="dense_1", keep_prob=keep_prob)
            net = tf_helper.dense_layer(net, self._conf.semi_n_hidden,
                                        name="dense_2", keep_prob=keep_prob)
            net = tf.nn.sigmoid(tf_helper.dense(net, flat_data_length,
                                                name="dense_3"))
        return net

    def _tf_semi_z_discriminator(self, style, keep_prob):
        """Discriminator for style codes.
        """
        with tf.variable_scope("semi_z_discriminator", reuse=tf.AUTO_REUSE):
            net = tf_helper.dense_layer(style, self._conf.semi_n_hidden,
                                        name="dense_1", keep_prob=keep_prob)
            net = tf_helper.dense_layer(net, self._conf.semi_n_hidden,
                                        name="dense_2", keep_prob=keep_prob)
            logits = tf_helper.dense(net, 1, name="dense_3")
        return logits

    def _tf_semi_y_discriminator(self, label, keep_prob):
        """Discriminator for class labels.
        """
        with tf.variable_scope("semi_y_discriminator", reuse=tf.AUTO_REUSE):
            net = tf_helper.dense_layer(label, self._conf.semi_n_hidden,
                                        name="dense_1", keep_prob=keep_prob)
            net = tf_helper.dense_layer(net, self._conf.semi_n_hidden,
                                        name="dense_2", keep_prob=keep_prob)
            logits = tf_helper.dense(net, 1, name="dense_3")
        return logits

    def _tf_define_optimizers(self, variables) -> None:
        """Prepare the training process. Define optimizers.
        """
        super()._tf_define_optimizers(variables)

        # FIXME[semi]: z and y discriminator
        var_z_discriminator = [var for var in variables
                               if "z_discriminator" in var.name]
        var_y_discriminator = [var for var in variables
                               if "y_discriminator" in var.name]
        var_generator = [var for var in variables
                         if "encoder" in var.name]

        self._op_z_discriminator = \
            tf.train.AdamOptimizer(learning_rate=self._tf_learning_rate/5).\
            minimize(self._tf_loss_discriminator,
                     global_step=self._tf_global_step,
                     var_list=var_z_discriminator)
        self._op_y_discriminator = \
            tf.train.AdamOptimizer(learning_rate=self._tf_learning_rate/5).\
            minimize(self._tf_loss_discriminator,
                     global_step=self._tf_global_step,
                     var_list=var_y_discriminator)
        self._op_generator = \
            tf.train.AdamOptimizer(learning_rate=self._tf_learning_rate).\
            minimize(self._tf_loss_generator,
                     global_step=self._tf_global_step,
                     var_list=var_generator)

        # optimizer for supervised data: minimize cross-entropy between
        # real and fake labels.
        self._op_crossentropy_labels = \
            tf.train.AdamOptimizer(learning_rate=self._tf_learning_rate).\
            minimize(self._crossentropy_labels,
                     global_step=self._tf_global_step,
                     var_list=var_generator)

    def _tf_train_step(self, sess, feed_dict) -> None:
        """Perform one training step.
        """
        # AutoEncoder phase
        super()._tf_train_step(sess, feed_dict)

        # Discriminator phase
        # FIXME[semi]: optimize both discriminators
        self._l_z_discriminator, _ = \
            sess.run([self._tf_loss_discriminator_style,
                      self._op_z_discriminator],
                     feed_dict=feed_dict)
        self._l_y_discriminator, _ = \
            sess.run([self._tf_loss_discriminator_label,
                      self._op_y_discriminator],
                     feed_dict=feed_dict)

        # Generator phase
        self._l_generator, _ = \
            sess.run([self._tf_loss_generator, self._op_generator],
                     feed_dict=feed_dict)

    def _tf_train_supervised_step(self, sess, feed_dict) -> None:
        # Cross_Entropy phase
        self._crossentropy, _ = \
            sess.run([self._crossentropy_labels,
                      self._op_crossentropy_labels],
                     feed_dict=feed_dict)

    def train(self,
              unlabeled_train_data: Datasource,
              labeled_train_data: Datasource,
              labeled_display_data: Datasource,
              sess: tf.Session, saver) -> None:
        """Train the network.
        """
        # display_data: (images, images_noised, labels)
        # A batch of data used for plotting intermediate results
        # during training (taken from the validation dataset)
        # Each is a numpy array of length 100 and appropriate shape.
        display_data = labeled_display_data[:100, ('array', 'array', 'label')]

        # obtain the variables used in the model
        total_vars = tf.trainable_variables()

        # Optimizers
        self._tf_define_optimizers(total_vars)

        #
        # Start the training
        #
        start_epoch = self._tf_initialize_variables(sess, saver)
        start_time = time.time()

        total_batch = len(labeled_train_data) // self._conf.batch_size
        labeled_batch_iterator = \
            labeled_train_data(batch_size=self._conf.batch_size,
                               attributes=('array', 'array', 'label'))
        unlabeled_batch_iterator = \
            unlabeled_train_data(batch_size=self._conf.batch_size,
                                 attributes=('array', 'noisy'))

        for epoch in tqdm(range(start_epoch, self._conf.n_epoch),
                          initial=start_epoch, total=self._conf.n_epoch):
            likelihood = 0
            discriminator_z_value = 0
            discriminator_y_value = 0
            generator_value = 0
            crossentropy_value = 0

            # Adapt the learning rate depending on the epoch
            lr_value = self._learning_rate_schedule(epoch)

            for _batch_idx in tqdm(range(total_batch)):

                #
                # Part 1: unsupervised learning
                #
                batch_xs, batch_noised_xs = next(unlabeled_batch_iterator)

                # FIXME[semi]:
                real_cat_labels = \
                    np.random.randint(low=0, high=self.nclasses,
                                      size=self._conf.batch_size)
                real_cat_labels = np.eye(self.nclasses)[real_cat_labels]

                # Sample from the prior distribution
                prior_style, _tf_prior_label_onehot = \
                    self._sample_prior(self._conf.prior, zdim=self.code_dim,
                                       nclasses=self.nclasses,
                                       batch_size=self._conf.batch_size,
                                       use_label_info=False)

                feed_dict = {
                    self._tf_inputs: batch_noised_xs,
                    self._tf_outputs: batch_xs,
                    # self._tf_labels: batch_ys,
                    self._tf_prior_style: prior_style,
                    # FIXME[semi]: real_cat_labels instead of
                    # prior_label_onehot
                    self._tf_prior_label: real_cat_labels,
                    self._tf_learning_rate: lr_value,
                    self._tf_keep_prob: self._conf.keep_prob
                }

                # perform the actual training
                self._tf_train_step(sess, feed_dict)

                #
                # Part 2: supervised training with labels
                #

                # Obtain a batch of labeled validation data to train
                # the encoder to predict correct class labels
                # (minimizing crossentropy)
                batch_semi_xs, batch_noised_semi_xs, batch_semi_ys = \
                    next(labeled_batch_iterator)

                feed_dict_semi = {
                    self._tf_inputs: batch_noised_semi_xs,
                    self._tf_outputs: batch_semi_xs,
                    self._tf_labels: batch_semi_ys,
                    # FIXME[semi]: _tf_prior_label was Y_cat
                    self._tf_prior_label: real_cat_labels,
                    self._tf_learning_rate: lr_value,
                    self._tf_keep_prob: self._conf.keep_prob
                }

                self._tf_train_supervised_step(self, sess, feed_dict_semi)

                # Summary
                likelihood += \
                    self._loss_reconstruction_value/total_batch
                discriminator_z_value += self._l_z_discriminator/total_batch
                discriminator_y_value += self._l_y_discriminator/total_batch
                generator_value += self._l_generator/total_batch
                crossentropy_value += self._crossentropy/total_batch

            # every 5th epoch (except the last) plot the manifold canvas
            if epoch % 5 == 0 or epoch == (self._conf.n_epoch - 1):
                name = f"Manifold_semi_canvas_{epoch}"
                self.plot_recoded_images(display_data[1],
                                         targets=display_data[0],
                                         filename=name)

            # output end of epoch information
            runtime = time.time() - start_time
            print(f"Epoch: {epoch:3d}, "
                  f"global step: {sess.run(self._tf_global_step)}, "
                  f"Time: {datetime.timedelta(seconds=runtime)}")
            print(f"             lr_AE: {lr_value:.5f}"
                  f"   loss_AE: {likelihood:.4f}")
            print(f"             lr_D: {lr_value/5:.5f}"
                  f"   loss_z_D: {discriminator_z_value:.4f},"
                  f"   loss_y_D: {discriminator_y_value:.4f}")
            print(f"             lr_G: {lr_value:.5f}"
                  f"   loss_G: {generator_value:.4f},"
                  f"   loss_CE: {crossentropy_value:.4f}\n")

def main() -> None:
    """The main program.
    """
    ModelClass = SupervisedAdversarialAutoencoder
    # ModelClass = SemiSupervisedAdversarialAutoencoder

    datasource_train = Datasource(module='mnist', one_hot=True)
    datasource_train.add_postprocessor(add_noise)
    datasource_test = Datasource(module='mnist', section='test', one_hot=True)

    aae = ModelClass(shape=datasource_train.shape, code_dim=2,
                     nclasses=len(datasource_train.label_scheme))
    aae.prepare()

    saver = tf.train.Saver(name='aae', filename="my_test_model")
    # print(f"\n\n##### Saver: {saver.last_checkpoints}\n\n")
    # print("tf.train.latest_checkpoint(): "
    #       f"{tf.train.latest_checkpoint('checkpoints')}")
    # print(f"\n\n##### Saver: {saver.last_checkpoints}\n\n")

    with tf.Session() as sess:
        aae.set_tensorflow_session(sess)
        aae.train(datasource_train, datasource_test, sess, saver)

        if aae.code_dim == 2:
            print("-" * 80)
            print("plot 2D Scatter Result")
            aae.plot_data_codes_2d(datasource_test._array,
                                   labels=datasource_test._labels,
                                   filename='2D_latent_space.png')

        if aae.code_dim and aae.conf.flag_plot_mlr:
            print("-" * 80)
            print("plot Manifold Learning Result")
            # filename = "PMLR/PMLR"
            aae.plot_decoded_codespace_2d(labels=5)

        if aae.conf.flag_plot_arr:
            print("-"*80)
            print("plot analogical reasoning result")
            aae.plot_analogical_decoding()

if __name__ == '__main__':
    main()


"""
针对brats数据集
做包含预处理的数据管道(Python生成器)

每次优先读取npy 不存在则读取nii 同时保存npy

读取240*240 指定155中部切片
或者读取240*240*155 3D

以2D为例，进行中心附近剪裁，单个240*240切片 中心有效区域为196*144
在这个区域内选取若然随机的128*128切片 进行训练 确保训练的数据是有监督的 对其的
切片范围为中心的196-128 和144-128  即 68*16范围内去随机值

"""

import os
import sys
from PIL import Image
import numpy as np
import nibabel as nib
from scipy import ndimage
import random
base = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(base, '../'))
from utils import CutPadding
import random
class UnpairedError(Exception):
    def __init__(self,path):
        self.err_msg = "There are not exixting paired samples! We can only find:"
        self.filename = path
class DataPipeLine():
    def __init__(self,path,target_size,patch_size,remake_flag=False,random_flag=False,crop="crop_random"):
        self.path = path
        self.datalist = self.__readDirFile(self.path,random_flag)
        self.target_size = target_size
        self.patch_size = patch_size
        self.dims = len(target_size)
        self.remake_flag = remake_flag
        self.random_flag = random_flag
        self.crop = crop.lower()
    def __readDirFile(self,path,random_flag=False):
        buf_A = []
        buf_B = []
        buf_A_mask_v0 = []
        buf_B_mask_v0 = []
        for (dirName, subdirList, fileList) in os.walk(path):
            try:
                for filename in fileList:
                    if "t1.nii" in filename.lower():
                        buf_A.append(os.path.join(dirName,filename))
                    if "t2.nii" in filename.lower():
                        buf_B.append(os.path.join(dirName,filename))
                    if "mask_t1_v0.nii" in filename.lower():
                        buf_A_mask_v0.append(os.path.join(dirName,filename))
                    if "mask_t2_v0.nii" in filename.lower():
                        buf_B_mask_v0.append(os.path.join(dirName,filename))
                if len(buf_A) > len(buf_B):
                    raise UnpairedError(buf_A.pop(-1))
                elif len(buf_A) < len(buf_B):
                    raise UnpairedError(buf_B.pop(-1))
                else:
                    pass
            except UnpairedError as error:
                print(error.err_msg)
                print(error.filename)
            else:# normal condition
                pass
            finally:# any way
                pass
        if random_flag:
            """
            打乱A B之间的对应关系
            """
            random_num1 = random.randint(0,200)
            random_num2 = random.randint(0,200)
            random.seed(random_num1)
            random.shuffle(buf_A)
            random.seed(random_num1)
            random.shuffle(buf_A_mask_v0)
            random.seed(random_num2)
            random.shuffle(buf_B)
            random.seed(random_num2)
            random.shuffle(buf_B_mask_v0)
            return list(zip(buf_A,buf_B,buf_A_mask_v0,buf_B_mask_v0))
        else:
            return list(zip(buf_A,buf_B,buf_A_mask_v0,buf_B_mask_v0))
    def re_rand(self):
        self.datalist = self.__readDirFile(self.path,random_flag=True)
    def read_file(self,path):
        if self.dims == 3:
            temp_path = path[:-3]+"npy"
            if (os.path.exists(temp_path)==True)and(self.remake_flag==False):
                return np.load(temp_path)
            else:
                return self.load_nii_file(path)
        elif self.dims == 2:
            temp_path = path[:-3]+"2D.npy"
            if (os.path.exists(temp_path)==True)and(self.remake_flag==False):
                return np.load(temp_path)
            else:
                return self.load_nii_file(path)
        else:
            raise ValueError
    def __read_nii_file(self,path):
        img = nib.load(path)
        img = np.array(img.dataobj[:,:,:])
        return img
    def __cut_nii_file(self,img):
        return CutPadding.cut_img_3D(img)
    def __save_nii2npy(self,img,path):
        if self.dims == 3:
            temp_path = path[:-3]+"npy"
        elif self.dims ==2:
            temp_path = path[:-3]+"2D.npy"
        else:
            raise ValueError
        np.save(temp_path,img)
        return img
    def __cut_np_array(self,array,target_shape=[128,128,128]):
        old_shape = array.shape
        buf = [0,0,0]
        for i in range(3):
            buf[i]=old_shape[i]//2-target_shape[i]//2
            #左半部右下标+1 减去目标点数的一半 获得新的起始点 10//2 -6//2 = 2 从下标2开始然后到下标2+6-1结束
        return array[buf[0]:buf[0]+target_shape[0],buf[1]:buf[1]+target_shape[1],buf[2]:buf[2]+target_shape[2]]
    def __normalize(self,slice,dtype=np.float32):
        tmp = slice/slice.max()
        return tmp.astype(dtype)
    def get_centro_ranges(self,target_size,patch_size):
        ranges_buf = []
        shape = target_size
        for i in range(len(patch_size)):
            if shape[i]<patch_size[i]:
                raise ValueError("Unsupported target size")
            elif shape[i]==patch_size[i]:
                pass
            else:
                diff = shape[i]-patch_size[i]
                begin = diff//2
                end = diff-begin
                ranges_buf.append([begin,shape[i]-end])
        return ranges_buf
    def load_nii_file(self,path):
        img = self.__read_nii_file(path)#读取3D源文件
        # img = self.__cut_nii_file(img)#去除文件周围无意义的区域 3D去黑边
        #缩放到目标大小 最近邻插值
        if len(self.target_size)==2:
            # temp_targer_size = self.target_size[:]+[self.target_size[-1]]
            temp_targer_size = self.target_size[:]+[155]
        else:
            temp_targer_size = self.target_size[:]
        # ratio = [temp_targer_size[x]/img.shape[x] for x in range(3)]
        # resize_image = ndimage.interpolation.zoom(img,ratio, mode='nearest')
        # assert resize_image.shape==tuple(temp_targer_size)
        # resize_image[resize_image<0]=0#去除插值后出现的负像素
        resize_image = CutPadding.center_crop_3D(img=img,target_size=temp_targer_size)
        if self.dims == 3:
            resize_image = resize_image
        elif self.dims ==2:
            resize_image = resize_image[:,:,temp_targer_size[-1]//2]
        else:
            raise ValueError
        img_norm = self.__normalize(resize_image,dtype=np.float32)#归一化
        img_saved = self.__save_nii2npy(img_norm,path)#保存 并且返回保存的文件 将对2D 3D区别对待
        return img_saved
    def __iter__(self):
        #实现__iter__ 本身就是一个迭代器 但是没有call方法 不能被tensorflow from_generator识别 所以必须在实现一个一般的生成器函数
        length = len(self.datalist)
        for i,(A,B,A_m0,B_m0) in enumerate(self.datalist):
            imgA = self.read_file(A)
            imgB = self.read_file(B)
            imgA_m0 = self.read_file(A_m0)
            imgB_m0 = self.read_file(B_m0)
            if self.dims == 3:
                buf = None
                yield (imgA,imgB,imgA_m0,imgB_m0,buf)
            elif self.dims == 2:
                buf = None
                yield (imgA,imgB,imgA_m0,imgB_m0,buf)
            else:
                raise ValueError("Unsupported dims")
            if (i+1)==length:
                if self.random_flag:
                    self.re_rand()
                    print("lueluelue")
        return None
    def generator(self):
        length = len(self.datalist)
        for i,(A,B,A_m0,B_m0) in enumerate(self.datalist):
            imgA = self.read_file(A)
            imgB = self.read_file(B)
            imgA_m0 = self.read_file(A_m0)
            imgB_m0 = self.read_file(B_m0)
            if self.dims == 3:
                ranges_buf = self.get_centro_ranges(target_size=self.target_size,patch_size=self.patch_size)
                slice_begin = ranges_buf[-1][0]
                slice_end = ranges_buf[-1][1]
                for slice_index in range(slice_begin,slice_end):
                    slice_imgA=imgA[:,:,slice_index]
                    slice_imgB=imgB[:,:,slice_index]
                    slice_imgA_m0=imgA_m0[:,:,slice_index]
                    slice_imgB_m0=imgB_m0[:,:,slice_index]
                    buf = np.array(ranges_buf[0:2],dtype=np.int32)
                    yield (slice_imgA,slice_imgB,slice_imgA_m0,slice_imgB_m0,buf)
            elif self.dims == 2:
                raise ValueError("Unsupported 2 dims")
            else:
                raise ValueError("Unsupported dims")
            if (i+1)==28:
                if self.random_flag:
                    self.re_rand()
                    print("lueluelue")
                break

        return None
    def chenk_saved_npy(self):
        #该方法直接进行一次全部迭代，将nii文件读取并且内容保存为预处理后的numpy矩阵 npy无压缩格式
        for i,(A,B,A_m0,B_m0,buf) in enumerate(self):
            print(i+1,A.shape,B.dtype,
                    B_m0.shape,A_m0.dtype,
                    A.max(),B.min(),
                    A.max(),B.min())
    def save_png(self):
        from PIL import Image
        for i,(A,B) in enumerate(self.datalist):
            imgA = np.array(255*self.read_file(A),dtype=np.uint8)
            imgB = np.array(255*self.read_file(B),dtype=np.uint8)
            imgA = Image.fromarray(imgA)
            imgB = Image.fromarray(imgB)
            imgA.save(A[:-4]+".png")
            imgB.save(B[:-4]+".png")
            print(i)
        return
if __name__ == "__main__":
    import tensorflow as tf
    a = DataPipeLine(path="G:\\Datasets\\BraTS\\ToCrop\\",
                     target_size=[240,240,155],
                     patch_size=[128,128,101],
                     remake_flag=False,
                     random_flag=True)
    # a.chenk_saved_npy()

    import matplotlib.pyplot as plt
    for i,(A,B,A_m0,B_m0,buf) in enumerate(a.generator()):
        if i == 10:
            plt.figure(figsize=(5,5))#图片大一点才可以承载像素
            plt.subplot(2,2,1)
            plt.imshow(A,cmap='gray')
            plt.axis('off')
            plt.subplot(2,2,2)
            plt.imshow(A_m0,cmap='gray')
            plt.axis('off')
            plt.subplot(2,2,3)
            plt.imshow(B,cmap='gray')
            plt.axis('off')
            plt.subplot(2,2,4)
            plt.imshow(B_m0,cmap='gray')
            plt.axis('off')
            print(buf)
            plt.show()
    for i,(A,B,A_m0,B_m0,buf) in enumerate(a.generator()):
        if i == 10:
            plt.figure(figsize=(5,5))#图片大一点才可以承载像素
            plt.subplot(2,2,1)
            plt.imshow(A,cmap='gray')
            plt.axis('off')
            plt.subplot(2,2,2)
            plt.imshow(A_m0,cmap='gray')
            plt.axis('off')
            plt.subplot(2,2,3)
            plt.imshow(B,cmap='gray')
            plt.axis('off')
            plt.subplot(2,2,4)
            plt.imshow(B_m0,cmap='gray')
            plt.axis('off')
            print(buf)
            plt.show()


# Shree KRISHNAya Namaha
# Warps frame using pose info
# Author: Nagabhushan S N
# Last Modified: 27/09/2021

import datetime
import time
import traceback
from pathlib import Path
from typing import Tuple, Optional

import numpy
import skimage.io

class Warper:
    def __init__(self, resolution: tuple = None):
        self.resolution = resolution
        return

    def forward_warp(self, frame1: numpy.ndarray, mask1: Optional[numpy.ndarray], depth1: numpy.ndarray,
                     transformation1: numpy.ndarray, transformation2: numpy.ndarray, intrinsic1: numpy.ndarray,
                     intrinsic2: Optional[numpy.ndarray]) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray,
                                                                   numpy.ndarray]:
        """
        Given a frame1 and global transformations transformation1 and transformation2, warps frame1 to next view using
        bilinear splatting.
        :param frame1: (h, w, 3) uint8 numpy array
        :param mask1: (h, w) bool numpy array. Wherever mask1 is False, those pixels are ignored while warping. Optional
        :param depth1: (h, w) float numpy array.
        :param transformation1: (4, 4) extrinsic transformation matrix of first view: [R, t; 0, 1]
        :param transformation2: (4, 4) extrinsic transformation matrix of second view: [R, t; 0, 1]
        :param intrinsic1: (3, 3) camera intrinsic matrix
        :param intrinsic2: (3, 3) camera intrinsic matrix. Optional
        """
        if self.resolution is not None:
            assert frame1.shape[:2] == self.resolution
        h, w = frame1.shape[:2]
        if mask1 is None:
            mask1 = numpy.ones(shape=(h, w), dtype=bool)
        if intrinsic2 is None:
            intrinsic2 = numpy.copy(intrinsic1)
        assert frame1.shape == (h, w, 3)
        assert mask1.shape == (h, w)
        assert depth1.shape == (h, w)
        assert transformation1.shape == (4, 4)
        assert transformation2.shape == (4, 4)
        assert intrinsic1.shape == (3, 3)
        assert intrinsic2.shape == (3, 3)

        trans_points1 = self.compute_transformed_points(depth1, transformation1, transformation2, intrinsic1,
                                                        intrinsic2)
        trans_coordinates = trans_points1[:, :, :2, 0] / trans_points1[:, :, 2:3, 0]
        trans_depth1 = trans_points1[:, :, 2, 0]

        grid = self.create_grid(h, w)
        flow12 = trans_coordinates - grid

        warped_frame2, mask2 = self.bilinear_splatting(frame1, mask1, trans_depth1, flow12, None, is_image=True)
        warped_depth2 = self.bilinear_splatting(trans_depth1[:, :, None], mask1, trans_depth1, flow12, None,
                                                is_image=False)[0][:, :, 0]
        return warped_frame2, mask2, warped_depth2, flow12

    def compute_transformed_points(self, depth1: numpy.ndarray, transformation1: numpy.ndarray,
                                   transformation2: numpy.ndarray, intrinsic1: numpy.ndarray,
                                   intrinsic2: Optional[numpy.ndarray]):
        """
        Computes transformed position for each pixel location
        """
        if self.resolution is not None:
            assert depth1.shape == self.resolution
        h, w = depth1.shape
        if intrinsic2 is None:
            intrinsic2 = numpy.copy(intrinsic1)
        transformation = numpy.matmul(transformation2, numpy.linalg.inv(transformation1))

        y1d = numpy.array(range(h))
        x1d = numpy.array(range(w))
        x2d, y2d = numpy.meshgrid(x1d, y1d)
        ones_2d = numpy.ones(shape=(h, w))
        ones_4d = ones_2d[:, :, None, None]
        pos_vectors_homo = numpy.stack([x2d, y2d, ones_2d], axis=2)[:, :, :, None]

        intrinsic1_inv = numpy.linalg.inv(intrinsic1)
        intrinsic1_inv_4d = intrinsic1_inv[None, None]
        intrinsic2_4d = intrinsic2[None, None]
        depth_4d = depth1[:, :, None, None]
        trans_4d = transformation[None, None]

        unnormalized_pos = numpy.matmul(intrinsic1_inv_4d, pos_vectors_homo)
        world_points = depth_4d * unnormalized_pos
        world_points_homo = numpy.concatenate([world_points, ones_4d], axis=2)
        trans_world_homo = numpy.matmul(trans_4d, world_points_homo)
        trans_world = trans_world_homo[:, :, :3]
        trans_norm_points = numpy.matmul(intrinsic2_4d, trans_world)
        return trans_norm_points

    def bilinear_splatting(self, frame1: numpy.ndarray, mask1: Optional[numpy.ndarray], depth1: numpy.ndarray,
                           flow12: numpy.ndarray, flow12_mask: Optional[numpy.ndarray], is_image: bool = False) -> \
            Tuple[numpy.ndarray, numpy.ndarray]:
        """
        Using inverse bilinear interpolation based splatting
        :param frame1: (h, w, c)
        :param mask1: (h, w): True if known and False if unknown. Optional
        :param depth1: (h, w)
        :param flow12: (h, w, 2)
        :param flow12_mask: (h, w): True if valid and False if invalid. Optional
        :param is_image: If true, the return array will be clipped to be in the range [0, 255] and type-casted to uint8
        :return: warped_frame2: (h, w, c)
                 mask2: (h, w): True if known and False if unknown
        """
        if self.resolution is not None:
            assert frame1.shape[:2] == self.resolution
        h, w, c = frame1.shape
        if mask1 is None:
            mask1 = numpy.ones(shape=(h, w), dtype=bool)
        if flow12_mask is None:
            flow12_mask = numpy.ones(shape=(h, w), dtype=bool)
        grid = self.create_grid(h, w)
        trans_pos = flow12 + grid

        trans_pos_offset = trans_pos + 1
        trans_pos_floor = numpy.floor(trans_pos_offset).astype('int')
        trans_pos_ceil = numpy.ceil(trans_pos_offset).astype('int')
        trans_pos_offset[:, :, 0] = numpy.clip(trans_pos_offset[:, :, 0], a_min=0, a_max=w + 1)
        trans_pos_offset[:, :, 1] = numpy.clip(trans_pos_offset[:, :, 1], a_min=0, a_max=h + 1)
        trans_pos_floor[:, :, 0] = numpy.clip(trans_pos_floor[:, :, 0], a_min=0, a_max=w + 1)
        trans_pos_floor[:, :, 1] = numpy.clip(trans_pos_floor[:, :, 1], a_min=0, a_max=h + 1)
        trans_pos_ceil[:, :, 0] = numpy.clip(trans_pos_ceil[:, :, 0], a_min=0, a_max=w + 1)
        trans_pos_ceil[:, :, 1] = numpy.clip(trans_pos_ceil[:, :, 1], a_min=0, a_max=h + 1)

        prox_weight_nw = (1 - (trans_pos_offset[:, :, 1] - trans_pos_floor[:, :, 1])) * \
                         (1 - (trans_pos_offset[:, :, 0] - trans_pos_floor[:, :, 0]))
        prox_weight_sw = (1 - (trans_pos_ceil[:, :, 1] - trans_pos_offset[:, :, 1])) * \
                         (1 - (trans_pos_offset[:, :, 0] - trans_pos_floor[:, :, 0]))
        prox_weight_ne = (1 - (trans_pos_offset[:, :, 1] - trans_pos_floor[:, :, 1])) * \
                         (1 - (trans_pos_ceil[:, :, 0] - trans_pos_offset[:, :, 0]))
        prox_weight_se = (1 - (trans_pos_ceil[:, :, 1] - trans_pos_offset[:, :, 1])) * \
                         (1 - (trans_pos_ceil[:, :, 0] - trans_pos_offset[:, :, 0]))

        sat_depth1 = numpy.clip(depth1, a_min=0, a_max=1000)
        log_depth1 = numpy.log(1 + sat_depth1)
        depth_weights = numpy.exp(log_depth1 / log_depth1.max() * 50)

        weight_nw = prox_weight_nw * mask1 * flow12_mask / depth_weights
        weight_sw = prox_weight_sw * mask1 * flow12_mask / depth_weights
        weight_ne = prox_weight_ne * mask1 * flow12_mask / depth_weights
        weight_se = prox_weight_se * mask1 * flow12_mask / depth_weights

        weight_nw_3d = weight_nw[:, :, None]
        weight_sw_3d = weight_sw[:, :, None]
        weight_ne_3d = weight_ne[:, :, None]
        weight_se_3d = weight_se[:, :, None]

        warped_image = numpy.zeros(shape=(h + 2, w + 2, c), dtype=numpy.float64)
        warped_weights = numpy.zeros(shape=(h + 2, w + 2), dtype=numpy.float64)

        numpy.add.at(warped_image, (trans_pos_floor[:, :, 1], trans_pos_floor[:, :, 0]), frame1 * weight_nw_3d)
        numpy.add.at(warped_image, (trans_pos_ceil[:, :, 1], trans_pos_floor[:, :, 0]), frame1 * weight_sw_3d)
        numpy.add.at(warped_image, (trans_pos_floor[:, :, 1], trans_pos_ceil[:, :, 0]), frame1 * weight_ne_3d)
        numpy.add.at(warped_image, (trans_pos_ceil[:, :, 1], trans_pos_ceil[:, :, 0]), frame1 * weight_se_3d)

        numpy.add.at(warped_weights, (trans_pos_floor[:, :, 1], trans_pos_floor[:, :, 0]), weight_nw)
        numpy.add.at(warped_weights, (trans_pos_ceil[:, :, 1], trans_pos_floor[:, :, 0]), weight_sw)
        numpy.add.at(warped_weights, (trans_pos_floor[:, :, 1], trans_pos_ceil[:, :, 0]), weight_ne)
        numpy.add.at(warped_weights, (trans_pos_ceil[:, :, 1], trans_pos_ceil[:, :, 0]), weight_se)

        cropped_warped_image = warped_image[1:-1, 1:-1]
        cropped_weights = warped_weights[1:-1, 1:-1]

        mask = cropped_weights > 0
        with numpy.errstate(invalid='ignore'):
            warped_frame2 = numpy.where(mask[:, :, None], cropped_warped_image / cropped_weights[:, :, None], 0)

        if is_image:
            assert numpy.min(warped_frame2) >= 0
            assert numpy.max(warped_frame2) <= 256
            clipped_image = numpy.clip(warped_frame2, a_min=0, a_max=255)
            warped_frame2 = numpy.round(clipped_image).astype('uint8')
        return warped_frame2, mask

    def bilinear_interpolation(self, frame2: numpy.ndarray, mask2: Optional[numpy.ndarray], flow12: numpy.ndarray,
                               flow12_mask: Optional[numpy.ndarray], is_image: bool = False) -> \
            Tuple[numpy.ndarray, numpy.ndarray]:
        """
        Using bilinear interpolation
        :param frame2: (h, w, c)
        :param mask2: (h, w): True if known and False if unknown. Optional
        :param flow12: (h, w, 2)
        :param flow12_mask: (h, w): True if valid and False if invalid. Optional
        :param is_image: If true, the return array will be clipped to be in the range [0, 255] and type-casted to uint8
        :return: warped_frame1: (h, w, c)
                 mask1: (h, w): True if known and False if unknown
        """
        if self.resolution is not None:
            assert frame2.shape[:2] == self.resolution
        h, w, c = frame2.shape
        if mask2 is None:
            mask2 = numpy.ones(shape=(h, w), dtype=bool)
        if flow12_mask is None:
            flow12_mask = numpy.ones(shape=(h, w), dtype=bool)
        grid = self.create_grid(h, w)
        trans_pos = flow12 + grid

        trans_pos_offset = trans_pos + 1
        trans_pos_floor = numpy.floor(trans_pos_offset).astype('int')
        trans_pos_ceil = numpy.ceil(trans_pos_offset).astype('int')
        trans_pos_offset[:, :, 0] = numpy.clip(trans_pos_offset[:, :, 0], a_min=0, a_max=w + 1)
        trans_pos_offset[:, :, 1] = numpy.clip(trans_pos_offset[:, :, 1], a_min=0, a_max=h + 1)
        trans_pos_floor[:, :, 0] = numpy.clip(trans_pos_floor[:, :, 0], a_min=0, a_max=w + 1)
        trans_pos_floor[:, :, 1] = numpy.clip(trans_pos_floor[:, :, 1], a_min=0, a_max=h + 1)
        trans_pos_ceil[:, :, 0] = numpy.clip(trans_pos_ceil[:, :, 0], a_min=0, a_max=w + 1)
        trans_pos_ceil[:, :, 1] = numpy.clip(trans_pos_ceil[:, :, 1], a_min=0, a_max=h + 1)

        prox_weight_nw = (1 - (trans_pos_offset[:, :, 1] - trans_pos_floor[:, :, 1])) * \
                         (1 - (trans_pos_offset[:, :, 0] - trans_pos_floor[:, :, 0]))
        prox_weight_sw = (1 - (trans_pos_ceil[:, :, 1] - trans_pos_offset[:, :, 1])) * \
                         (1 - (trans_pos_offset[:, :, 0] - trans_pos_floor[:, :, 0]))
        prox_weight_ne = (1 - (trans_pos_offset[:, :, 1] - trans_pos_floor[:, :, 1])) * \
                         (1 - (trans_pos_ceil[:, :, 0] - trans_pos_offset[:, :, 0]))
        prox_weight_se = (1 - (trans_pos_ceil[:, :, 1] - trans_pos_offset[:, :, 1])) * \
                         (1 - (trans_pos_ceil[:, :, 0] - trans_pos_offset[:, :, 0]))

        weight_nw = prox_weight_nw * flow12_mask
        weight_sw = prox_weight_sw * flow12_mask
        weight_ne = prox_weight_ne * flow12_mask
        weight_se = prox_weight_se * flow12_mask

        weight_nw_3d = weight_nw[:, :, None]
        weight_sw_3d = weight_sw[:, :, None]
        weight_ne_3d = weight_ne[:, :, None]
        weight_se_3d = weight_se[:, :, None]

        frame2_offset = numpy.pad(frame2, pad_width=((1, 1), (1, 1), (0, 0)), mode='constant', constant_values=0)
        mask2_offset = numpy.pad(mask2, pad_width=((1, 1), (1, 1)), mode='constant', constant_values=0)

        f2_nw = frame2_offset[trans_pos_floor[:, :, 1], trans_pos_floor[:, :, 0]]
        f2_sw = frame2_offset[trans_pos_ceil[:, :, 1], trans_pos_floor[:, :, 0]]
        f2_ne = frame2_offset[trans_pos_floor[:, :, 1], trans_pos_ceil[:, :, 0]]
        f2_se = frame2_offset[trans_pos_ceil[:, :, 1], trans_pos_ceil[:, :, 0]]

        m2_nw = mask2_offset[trans_pos_floor[:, :, 1], trans_pos_floor[:, :, 0]]
        m2_sw = mask2_offset[trans_pos_ceil[:, :, 1], trans_pos_floor[:, :, 0]]
        m2_ne = mask2_offset[trans_pos_floor[:, :, 1], trans_pos_ceil[:, :, 0]]
        m2_se = mask2_offset[trans_pos_ceil[:, :, 1], trans_pos_ceil[:, :, 0]]

        m2_nw_3d = m2_nw[:, :, None]
        m2_sw_3d = m2_sw[:, :, None]
        m2_ne_3d = m2_ne[:, :, None]
        m2_se_3d = m2_se[:, :, None]

        nr = weight_nw_3d * f2_nw * m2_nw_3d + weight_sw_3d * f2_sw * m2_sw_3d + \
             weight_ne_3d * f2_ne * m2_ne_3d + weight_se_3d * f2_se * m2_se_3d
        dr = weight_nw_3d * m2_nw_3d + weight_sw_3d * m2_sw_3d + weight_ne_3d * m2_ne_3d + weight_se_3d * m2_se_3d
        warped_frame1 = numpy.where(dr > 0, nr / dr, 0)
        mask1 = dr[:, :, 0] > 0

        if is_image:
            assert numpy.min(warped_frame1) >= 0
            assert numpy.max(warped_frame1) <= 256
            clipped_image = numpy.clip(warped_frame1, a_min=0, a_max=255)
            warped_frame1 = numpy.round(clipped_image).astype('uint8')
        return warped_frame1, mask1

    @staticmethod
    def create_grid(h, w):
        x_1d = numpy.arange(0, w)[None]
        y_1d = numpy.arange(0, h)[:, None]
        x_2d = numpy.repeat(x_1d, repeats=h, axis=0)
        y_2d = numpy.repeat(y_1d, repeats=w, axis=1)
        grid = numpy.stack([x_2d, y_2d], axis=2)
        return grid

    @staticmethod
    def read_image(path: Path) -> numpy.ndarray:
        if path.suffix in ['.jpg', '.png', '.bmp']:
            image = skimage.io.imread(path.as_posix())
        elif path.suffix == '.npy':
            image = numpy.load(path.as_posix())
        else:
            raise RuntimeError(f'Unknown image format: {path.as_posix()}')
        return image

    @staticmethod
    def read_depth(path: Path) -> numpy.ndarray:
        if path.suffix == '.png':
            depth = skimage.io.imread(path.as_posix())
        elif path.suffix == '.npy':
            depth = numpy.load(path.as_posix())
        elif path.suffix == '.npz':
            with numpy.load(path.as_posix()) as depth_data:
                depth = depth_data['depth']
        elif path.suffix == '.exr':
            import Imath
            import OpenEXR

            exr_file = OpenEXR.InputFile(path.as_posix())
            raw_bytes = exr_file.channel('B', Imath.PixelType(Imath.PixelType.FLOAT))
            depth_vector = numpy.frombuffer(raw_bytes, dtype=numpy.float32)
            height = exr_file.header()['displayWindow'].max.y + 1 - exr_file.header()['displayWindow'].min.y
            width = exr_file.header()['displayWindow'].max.x + 1 - exr_file.header()['displayWindow'].min.x
            depth = numpy.reshape(depth_vector, (height, width))
        else:
            raise RuntimeError(f'Unknown depth format: {path.as_posix()}')
        return depth

    @staticmethod
    def camera_intrinsic_transform(capture_width=1920, capture_height=1080, patch_start_point: tuple = (0, 0)):
        start_y, start_x = patch_start_point
        camera_intrinsics = numpy.eye(3)
        camera_intrinsics[0, 0] = 2100
        camera_intrinsics[0, 2] = capture_width / 2.0 - start_x
        camera_intrinsics[1, 1] = 2100
        camera_intrinsics[1, 2] = capture_height / 2.0 - start_y
        return camera_intrinsics

def demo1():
    frame1_path = Path('../Data/frame1.png')
    frame2_path = Path('../Data/frame2.png')
    depth1_path = Path('../Data/depth1.npy')
    transformation1 = numpy.array([
        4.067366123199462891e-01, 9.135454893112182617e-01, 2.251522164442576468e-05, -1.571802258491516113e+00,
        -7.961163669824600220e-02, 3.546993434429168701e-02, -9.961947202682495117e-01, 1.842712044715881348e+00,
        -9.100699424743652344e-01, 4.051870703697204590e-01, 8.715576678514480591e-02, -2.255212306976318359e+00,
        0.000000000000000000e+00, 0.000000000000000000e+00, 0.000000000000000000e+00, 1.000000000000000000e+00
    ]).reshape(4, 4)
    transformation2 = numpy.array([
        4.067366123199462891e-01, 9.135454893112182617e-01, 2.251522164442576468e-05, -1.616834521293640137e+00,
        -7.961163669824600220e-02, 3.546993434429168701e-02, -9.961947202682495117e-01, 1.848096847534179688e+00,
        -9.100699424743652344e-01, 4.051870703697204590e-01, 8.715576678514480591e-02, -2.275809526443481445e+00,
        0.000000000000000000e+00, 0.000000000000000000e+00, 0.000000000000000000e+00, 1.000000000000000000e+00
    ]).reshape(4, 4)

    warper = Warper()
    frame1 = warper.read_image(frame1_path)
    frame2 = warper.read_image(frame2_path)
    depth1 = warper.read_depth(depth1_path)
    intrinsic = warper.camera_intrinsic_transform()

    warped_frame2 = warper.forward_warp(frame1, None, depth1, transformation1, transformation2, intrinsic, None)[0]
    skimage.io.imsave('frame1.png', frame1)
    skimage.io.imsave('frame2.png', frame2)
    skimage.io.imsave('frame2_warped.png', warped_frame2)
    return

def main():
    demo1()
    return

if __name__ == '__main__':
    print('Program started at ' + datetime.datetime.now().strftime('%d/%m/%Y %I:%M:%S %p'))
    start_time = time.time()
    try:
        main()
    except Exception as e:
        print(e)
        traceback.print_exc()
    end_time = time.time()
    print('Program ended at ' + datetime.datetime.now().strftime('%d/%m/%Y %I:%M:%S %p'))
    print('Execution time: ' + str(datetime.timedelta(seconds=end_time - start_time)))


#!/usr/bin/env python

import sys
import loco
import tinymath as tm
import numpy as np

PHYSICS_BACKEND = loco.sim.PHYSICS_NONE
RENDERING_BACKEND = loco.sim.RENDERING_GLVIZ_GLFW

class DoublePendulum(object):

    def __init__(self, name='double_pendulum', pos=[0.0, 0.0, 2.0], l1=0.5, l2=0.5):
        super(DoublePendulum, self).__init__()

        self.kintree = loco.sim.kintree.KinematicTree(name, pos, tm.Matrix3f())
        self.base = loco.sim.kintree.Box(name + '_base', [0.1, 0.1, 0.1])
        self.link_1 = loco.sim.kintree.Capsule(name + '_link_1', 0.05, l1)
        self.link_2 = loco.sim.kintree.Capsule(name + '_link_2', 0.04, l2)

        jnt_1_axis = [1.0, 0.0, 0.0]
        jnt_2_axis = [1.0, 0.0, 0.0]
        self.jnt_base = loco.sim.kintree.FixedJoint(name + '_jnt_fixed')
        self.jnt_1 = loco.sim.kintree.RevoluteJoint(name + '_jnt_rev_1', jnt_1_axis, [1.0, -1.0])
        self.jnt_2 = loco.sim.kintree.RevoluteJoint(name + '_jnt_rev_2', jnt_2_axis, [1.0, -1.0])

        tf_link_1_to_base = tm.Matrix4f(np.identity(3), [0.0, 0.0, -0.5 * l1])
        tf_link_2_to_link_1 = tm.Matrix4f(np.identity(3), [0.0, 0.0, -0.5 * l1 - 0.5 * l2] )
        tf_jnt_1_to_link_1 = tm.Matrix4f(np.identity(3), [0.0, 0.0, 0.5 * l1])
        tf_jnt_2_to_link_2 = tm.Matrix4f(np.identity(3), [0.0, 0.0, 0.5 * l2])

        self.kintree.SetRoot(self.base)
        self.base.AddJoint(self.jnt_base, np.identity(4))
        self.base.AddChild(self.link_1, tf_link_1_to_base)
        self.link_1.AddJoint(self.jnt_1, tf_jnt_1_to_link_1)
        self.link_1.AddChild(self.link_2, tf_link_2_to_link_1)
        self.link_2.AddJoint(self.jnt_2, tf_jnt_2_to_link_2)

if __name__ == '__main__' :

    if len( sys.argv ) > 1 :
        choice_backend = sys.argv[1]
        if choice_backend == 'mujoco' :
            PHYSICS_BACKEND = loco.sim.PHYSICS_MUJOCO
        elif choice_backend == 'bullet' :
            PHYSICS_BACKEND = loco.sim.PHYSICS_BULLET
        elif choice_backend == 'dart' :
            PHYSICS_BACKEND = loco.sim.PHYSICS_DART
        elif choice_backend == 'raisim' :
            PHYSICS_BACKEND = loco.sim.PHYSICS_RAISIM
    print( 'Physics backend: {}'.format( PHYSICS_BACKEND ) )
    print( 'Rendering backend: {}'.format( RENDERING_BACKEND ) )

    scenario = loco.sim.Scenario()
    floor = scenario.AddSingleBody( loco.sim.primitives.Plane( "floor", 10.0, 10.0, tm.Vector3f(), tm.Matrix3f() ) )
    double_pendulum = DoublePendulum()
    scenario.AddKinematicTree( double_pendulum.kintree )

    runtime = loco.sim.Runtime( PHYSICS_BACKEND, RENDERING_BACKEND )
    simulation = runtime.CreateSimulation( scenario )
    visualizer = runtime.CreateVisualizer( scenario )

    floor.drawable.texture = 'built_in_chessboard'
    floor.drawable.ambient = [ 0.3, 0.5, 0.7 ]
    floor.drawable.diffuse = [ 0.3, 0.5, 0.7 ]
    floor.drawable.specular = [ 0.3, 0.5, 0.7 ]

    while visualizer.IsActive() :
        if visualizer.CheckSingleKeyPress( loco.sim.Keys.KEY_ESCAPE ) :
            break
        elif visualizer.CheckSingleKeyPress( loco.sim.Keys.KEY_R ) :
            simulation.Reset()
        elif visualizer.CheckSingleKeyPress( loco.sim.Keys.KEY_P ) :
            simulation.Pause() if simulation.running else simulation.Resume()

        double_pendulum.jnt_1.angle = double_pendulum.jnt_1.angle + 0.01
        double_pendulum.jnt_2.angle = double_pendulum.jnt_2.angle + 0.005

        simulation.Step( 1. / 60. )
        visualizer.Render()

    runtime.DestroySimulation()
    runtime.DestroyVisualizer()


"""
ゲルフォントの定数
"""

def gelfonds_constant(digit: int) -> str:
	from sympy import N, I
	return str(N("(-1) ** -I", digit))

if __name__ == "__main__":
	digit = 10000
	print(gelfonds_constant(digit))


"""Plot figure 6: attack rates vs R0."""
import numpy as np
import pandas as pd
from scipy.stats import linregress

import matplotlib as mplt
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.cm as cm
import matplotlib.font_manager as font_manager

import cmocean
import cmasher
import seaborn
import copy
import os

# set the font family style
mplt.rcParams['font.family'] = 'Myriad Pro'  # change to a font you have installed on your computer - checkout Google Fonts for free fonts available for download

# set some initial paths

# path to the directory where this script lives
thisdir = os.path.abspath('')

# path to the main directory of the repository
maindir = os.path.split(os.path.split(thisdir)[0])[0]

# path to the analysis_results subdirectory
analysisdir = os.path.split(thisdir)[0]

# path to the data subdirectory
datadir = os.path.join(os.path.split(os.path.split(thisdir)[0])[0], 'data')

# path to the figures subdirectory within analysis_results
figdir = os.path.join(analysisdir, 'figures')

location_file = os.path.join(analysisdir, 'location_country_names.csv')
locations_df = pd.read_csv(location_file, delimiter = ',')

setting_fractions_file = os.path.join(analysisdir, 'summary_fractions_by_location.csv')
setting_fractions_df = pd.read_csv(setting_fractions_file)
setting_codes = ['H','S','W','R']

def get_country_name(df,location):
    """
    Return country name of the location

    Args:
        df (pandas Dataframe) : dataframe containing locations and corresponding countries
        location (str)        : name of the location

    Returns:
        str: Name of the country the location is in.
    """
    d = df[df.location == location]
    return d.country.values[0]

def get_locations_by_country(df,country):
    """
    Return locations in the country

    Args:
        df (pandas Dataframe) : dataframe containing locations and corresponding countries
        country (str)         : name of the country

    Returns:
        str: Name of the country the location is in.
    """
    locations = list(df[df.country == country].location.values)
    return locations

def get_fractions(setting_fractions_df, location):
    """
    Get the fraction of people with contacts in each setting of household (H), school (S), or work (W).

    Args:
        setting_fractions_df (pandas DataFrame) : a dataframe
        location (str)                          : name of the location

    Returns:
        dict: A dictionary of fractions of people with contacts in each setting for the location.
    """
    fractions = dict.fromkeys(['H','S','W','R'],1.)
    d = setting_fractions_df[setting_fractions_df.location == location]
    fractions['H'] = d.NhN.values[0]
    fractions['S'] = d.NsN.values[0]
    fractions['W'] = d.NwN.values[0]
    return fractions

def read_contact_matrix(location, country, level, setting, num_agebrackets=85):
    """
    Read in the contact for each setting.

    Args:
        location (str)        : name of the location
        country (str)         : name of the country
        level (str)           : name of level (country or subnational)
        setting (str)         : name of the contact setting
        num_agebrackets (int) : the number of age brackets for the matrix

    Returns:
        A numpy matrix of contact.
    """
    setting_type, setting_suffix = 'F', 'setting'
    if setting == 'overall':
        setting_type, setting_suffix = 'M', 'contact_matrix'

    if country == 'Europe':
        country = location
        level = 'country'

    if level == 'country':
        file_name = country + '_' + level + '_level_' + setting_type + '_' + setting + '_' + setting_suffix + '_' + '%i' % num_agebrackets + '.csv'
    else:
        file_name = country + '_' + level + '_' + location + '_' + setting_type + '_' + setting + '_' + setting_suffix + '_' + '%i' % num_agebrackets + '.csv'
    file_path = os.path.join(datadir, 'contact_matrices', file_name)
    M = np.loadtxt(file_path, delimiter=',')
    return M

def get_ages(location, country, level, num_agebrackets=85):
    """
    Get the age count for the synthetic population of the location.

    Args:
        location (str)        : name of the location
        country (str)         : name of the country
        level (str)           : name of level (country or subnational)
        num_agebrackets (int) : the number of age brackets

    Returns:
        dict: A dictionary of the age count.
    """

    if country == 'Europe':
        country = location
        level = 'country'

    if level == 'country':
        file_name = country + '_' + level + '_level_age_distribution_' + '%i' % num_agebrackets + '.csv'
    else:
        file_name = country + '_' + level + '_' + location + '_age_distribution_' + '%i' % num_agebrackets + '.csv'
    file_path = os.path.join(datadir, 'age_distributions', file_name)
    df = pd.read_csv(file_path, delimiter=',', header=None)
    df.columns = ['age', 'age_count']
    ages = dict(zip(df.age.values.astype(int), df.age_count.values))
    return ages

def get_average_age(ages):
    """
    Get the average age from a dictionary of age counts.

    Args:
        ages (dict): dictionary of age counts

    Return:
        float: The average age given the age count.
    """
    average_age = 0
    total_population = sum(ages.values())
    for a in ages:
        average_age += ages[a] * a
    average_age = average_age/total_population
    return average_age

def get_school_age_distribution(location):
    """
    Get the age count of people active in the school setting.

    Args:
        location (str): name of the location

    Returns:
        dict: Age count of people active in the school setting.
    """
    ages = {}
    file_path = os.path.join(analysisdir, 'schools_age_distributions',
                'schools_age_distributions_' + location + '.dat')
    df = pd.read_csv(file_path, delimiter = ',')
    ages = dict(zip(df.age.values, df.setting_count.values))
    return ages

def get_percent_in_school(ages, school_ages):
    """
    Get the percent of people in school.

    Args:
        ages (dict)        : age count
        school_ages (dict) : school age count

    Returns:
        float: The percent of people in the school setting.
    """
    total_in_school = np.sum([v for v in school_ages.values()], dtype = float)
    total_population = np.sum([v for v in ages.values()], dtype = float)
    return total_in_school/total_population * 100

def get_attack_rates_df(reference_location, reference_scenario, beta, susceptibility_drop_factor, gamma_inverse, num_agebrackets):
    """
    Get attack rates dataframe for an SIR compartmental model with age specific contact patterns.

    Args:
        reference_location (str)           : name of reference location or locations
        reference_scenario (str)           : specific reference scenario
        beta (float)                       : the transmissibilty
        susceptibility_drop_factor (float) : susceptibility of adults to those under 18
        gamma_inverse (float)              : the mean recovery period
        num_agebrackets (int)              : the number of age brackets for the matrix

    Returns:
        Pandas dataframe of attack rates by location
    """

    file_path = os.path.join(analysisdir, 'reference_location_' + reference_location, 'mcmc_beta_and_dropfactor_scenario_' + reference_scenario,
                             'all_locations_attack_rates_by_age_reference_scenario_' + reference_scenario + '_beta_' + '%.2f' % beta + '_susceptibility_' + '%.2f' % susceptibility + '_gamma_inverse_' + '%.1f' % gamma_inverse + '_' + str(num_agebrackets) + '.csv')
    return pd.read_csv(file_path)

def get_attack_rate(df, location):
    """
    Get the total attack rate for the location from the dataframe for an SIR compartmental model with age specific contact patterns.

    Args:
        df (pd.DataFrame) : a dataframe of attack rates by location
        location (str)    : name of the location

    Returns:
        float: The total attack rate for the location as a fraction from an SIR compartmental model with age specific contact patterns. Values between 0 and 1.
    """
    return df.loc[df['location'] == location]['artotal'].values[0]

def get_homogeneous_attack_rate_df(gamma_inverse, num_agebrackets):
    """
    Get a dataframe with the attack rate for an SIR model with the homogeneous mixing assumption for different basic reproduction, R0, values with a given average recovery period.

    Args:
        gamma_inverse (float): the mean recovery period
        num_agebrackets (int): the number of age brackets for the matrix

    Returns:
        Pandas dataframe of attack rates by R0 value
    """
    file_path = os.path.join(analysisdir, 'homogeneous_sir_attack_rates', 'attack_rates_SIR_homogeneous_mixing.csv')
    df = pd.read_csv(file_path)
    return df

def get_homogeneous_attack_rate(df, R0_star):
    """
    Get the attack rate for an SIR model with the homogeneous mixing assumption for a given basic reproduction, R0_star, value.

    Args:
        df (pd.DataFrame) : a dataframe of attack rates by the basic reproduction number
        R0_star (float)   : the basic reproduction number

    Returns:
        float: The total attack rate as a fraction from an SIR model with homogeneous mixing assumptions and specified basic reproduction number R0_star. Values between 0 and 1.
    """
    return df.loc[df['R0'] == R0_star]['attack_rate'].values[0]

def get_eigenvalue(matrix):
    """
    Get the real component of the leading eigenvalue of a square matrix.

    Args:
        matrix (np.ndarray): square matrix

    Returns:
        float: Real component of the leading eigenvalue of the matrix.
    """
    eigenvalue = max(np.linalg.eigvals(matrix)).real
    return eigenvalue

def get_R0(beta, gamma_inverse, matrix):
    """
    Get the basic reproduction number, R0, for an SIR compartmental model given the basic reproduction number, the mean recovery period, and the contact matrix.

    Args:
        beta (float)          : the transmissibility beta
        gamma_inverse (float) : the mean recovery period
        matrix (np.ndarray)   : the contact matrix

    Returns:
        float: The basic reproduction number R0 for an SIR compartmental model.
    """

    gamma = 1./gamma_inverse
    eigenvalue = get_eigenvalue(matrix)
    return beta * eigenvalue / gamma

def get_beta(R0, gamma_inverse, matrix):
    """
    Get the transmissibility from an SIR model with age specific contact patterns and basic reproduction number.

    Args:
        R0_star (float)       : the basic reproduction number
        gamma_inverse (float) : the mean recovery period
        matrix (np.ndarray)   : the age specific contact matrix

    Returns:
        float: The transmissibility from an SIR model with age specific contact patterns.
    """
    gamma = float(1)/gamma_inverse
    eigenvalue = get_eigenvalue(matrix)
    return R0 * gamma / eigenvalue

def linear_function(x,m,b):
    """
    Get the y value of a linear function given the x value.

    Args:
        m (float): the slope
        b (float): the intercept

    Returns:
        The expected y value from a linear function at some specified x value.
    """
    return m*x + b

def plot_fig(countries, reference_location, reference_scenario, beta, susceptibility_drop_factor, gamma_inverse, num_agebrackets):
    """
    Plot the attack rates from an SIR model with age specific contact patterns vs the basic reproduction, the average age,
    and the percent of the population with contacts in the school layer for subnational locations.

    Args:
        countries (list)                   : list of countries
        reference_location (str)           : the reference location (or set of locations) the transmissibilty is calibrated to
        reference_scenario (str)           : label of the reference scenario
        beta (float)                       : the transmissibility
        susceptibility_drop_factor (float) : susceptibility of adults to those under 18
        gamma_inverse (float)              : the mean recovery period
        num_agebrackets (int)              : the number of age brackets

    Returns:
        Matplotlib figure.
    """
    countries = ['Australia', 'Canada', 'China', 'Europe', 'India', 'Israel', 'Japan', 'Russia', 'South_Africa', 'United_States']
    locations = []

    for country in countries:
        locations += get_locations_by_country(locations_df, country)
        if country in locations and country != 'Israel':
            locations.remove(country)

        if country == 'India':
            locations.remove('Dadra_and_Nagar_Haveli')
            locations.remove('Chandigarh')
            locations.remove('Lakshadweep')

    if 'Russian_Federation' in locations:
        locations.remove('Russian_Federation')

    country_colour_dic = {}
    country_colour_dic['Australia'] = '#0000ff'
    country_colour_dic['Canada'] = '#2ab207'
    country_colour_dic['China'] = '#fcc200'
    country_colour_dic['Europe'] = '#941cca'
    country_colour_dic['India'] = 'darkorange'
    country_colour_dic['Israel'] = '#9b9b9b'
    country_colour_dic['Japan'] = '#000098'
    country_colour_dic['Russia'] = '#dc142b'
    country_colour_dic['South_Africa'] = '#b5d93c'
    country_colour_dic['United_States'] = '#00ace7'

    hdf = get_homogeneous_attack_rate_df(gamma_inverse, num_agebrackets)
    df = get_attack_rates_df(reference_location, reference_scenario, beta, susceptibility_drop_factor, gamma_inverse, num_agebrackets)

    width = 16
    height = 5
    left = 0.06
    right = 0.865
    bottom = 0.16
    top = 0.88
    wspace = 0.32

    fig, ax = plt.subplots(1, 3, figsize=(width, height))
    fig.subplots_adjust(left = left, right = right, top = top, bottom = bottom, wspace = wspace)

    leg_left = right + 0.01
    leg_right = 0.985
    leg_bottom = bottom
    leg_top = top
    leg_width = leg_right - leg_left
    leg_height = leg_top - leg_bottom

    fontsize = 20

    axleg = fig.add_axes([leg_left, leg_bottom, leg_width, leg_height])
    axleg.axis('off')

    attack_rates_list = []
    R0_list = []
    average_age_list = []
    percent_in_school_list = []
    color_list = []

    beta_drop_age = 18

    for n, location in enumerate(locations):
        country = get_country_name(locations_df, location)
        if location == 'Israel':
            matrix = read_contact_matrix(location, country, 'country', 'overall')
            ages = get_ages(location, country, 'country')
        else:
            matrix = read_contact_matrix(location, country, 'subnational', 'overall')
            ages = get_ages(location, country, 'subnational')

        R0 = get_R0(beta, gamma_inverse, matrix)
        R0_list.append(R0)

        average_age = get_average_age(ages)
        average_age_list.append(average_age)

        if country != 'Europe':
            school_ages = get_school_age_distribution(location)
            percent_in_school = get_percent_in_school(ages, school_ages)

        else:
            fractions = get_fractions(setting_fractions_df, location)
            percent_in_school = fractions['S'] * 100

        percent_in_school_list.append(percent_in_school)

        ar = get_attack_rate(df, location) * 100
        attack_rates_list.append(ar)
        color = country_colour_dic[country]
        color_list.append(color)

    homogeneous_attack_rates_list = hdf.attack_rate.values * 100
    homogeneous_R0_list = hdf.R0.values

    size = 12

    leg = []

    ax[0].scatter(R0_list, attack_rates_list, marker='o', color=color_list, s=size)
    ax[0].plot(homogeneous_R0_list, homogeneous_attack_rates_list, color='k', lw=1.5, label='Homogenous \nmixing model')
    leg.append(ax[0].legend(loc=2, fontsize=17))
    ax[0].set_xlim(1.5, 2.0)
    ax[0].set_xticks(np.arange(1.5, 2.01, 0.1))
    ax[0].set_xlabel(r'$R_0$', fontsize = fontsize)
    ax[0].text(1.39, 50, 'a', fontsize=fontsize+24, fontstyle='oblique')

    m,b,r,p,std_err = linregress(average_age_list,attack_rates_list)
    y_theory = np.array([linear_function(a,m,b) for a in average_age_list])
    ax[1].plot(average_age_list, y_theory, color = 'k', lw = 1.5)
    ax[1].scatter(average_age_list, attack_rates_list, marker='o', color=color_list, s=size)
    ax[1].text(45,73, r'$\rho$ = ' + '%.2f' % r , fontsize = 18, verticalalignment = 'center', horizontalalignment = 'center', color = 'k')
    ax[1].set_xlim(20, 55)
    ax[1].set_xticks(np.arange(20, 51, 10))
    leg.append(ax[1].legend(loc = 7, fontsize = 16, ncol = 1))
    ax[1].set_xlabel('Average Age', fontsize = fontsize)
    ax[1].text(13.5, 50, 'b', fontsize=fontsize+24, fontstyle='oblique')

    m,b,r,p,std_err = linregress(percent_in_school_list,attack_rates_list)
    y_theory = np.array([linear_function(a,m,b) for a in percent_in_school_list])
    ax[2].plot(percent_in_school_list, y_theory, color = 'k', lw = 1.5)
    ax[2].scatter(percent_in_school_list, attack_rates_list, marker='o', color=color_list, s=size)
    ax[2].text(20,73, r'$\rho$ = ' + '%.2f' % r , fontsize = 18, verticalalignment = 'center', horizontalalignment = 'center', color = 'k')
    ax[2].set_xlim(10, 45)
    ax[2].set_xticks(np.arange(10, 50, 10))
    ax[2].set_xlabel('% In Educational Institutions', fontsize = fontsize)
    ax[2].text(3.5, 50, 'c', fontsize=fontsize+24, fontstyle='oblique')

    for country in ['Australia','Canada','China','Europe','India','Israel','Japan','Russia','South_Africa','United_States']:
        ax[2].scatter(0,0, color = country_colour_dic[country], s = size * 2, label = country.replace('-',' ').replace('_',' '))
    leg.append(ax[2].legend(loc = 7, fontsize = 17, ncol = 1, bbox_to_anchor = (0.6,0.4,1,0.2)))

    for i in range(len(ax)):
        ax[i].set_ylabel('Attack Rate (%)', fontsize = fontsize)
        ax[i].set_yticks(np.arange(55, 81, 5))
        ax[i].set_ylim(54, 80)
        ax[i].tick_params(labelsize=fontsize-2)
        ax[i].tick_params(axis='x', which='minor', bottom=False)
        ax[i].tick_params(axis='y', which='minor', left=False)
        leg[i].draw_frame(False)

    plt.minorticks_off()

    fig_path = os.path.join(figdir, 'fig_6.pdf')
    fig.savefig(fig_path, format='pdf')

if __name__ == '__main__':

    reference_location = 'polymod_and_tomsk'
    reference_scenario = 'all_locations'

    beta = 0.04752
    susceptibility = 1.0
    gamma_inverse = 2.6
    num_agebrackets = 85

    countries = ['Australia', 'Canada', 'China', 'Europe', 'India', 'Israel', 'Japan', 'Russia', 'South_Africa', 'United_States']

    plot_fig(countries, reference_location, reference_scenario, beta, susceptibility, gamma_inverse, num_agebrackets)