import math
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
import copy
import os

from collections import OrderedDict

from fairseq import utils, distributed_utils
from fairseq.models import (
    FairseqEncoderDecoderModel,
    BaseFairseqModel,
    register_model,
    register_model_architecture,
)
from fairseq.models.transformer import (
    base_architecture,
    Embedding,
    TransformerModel,
    FairseqEncoder,
    FairseqIncrementalDecoder,
    EncoderOut,
)

from fairseq.modules import (
    AdaptiveSoftmax,
    LayerNorm,
    PositionalEmbedding,
    SinusoidalPositionalEmbedding,
    TransformerDecoderLayer,
    TransformerEncoderLayer,
)

@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    try:
        tensors_gather = [torch.zeros_like(tensor)
            for _ in range(torch.distributed.get_world_size())]
        # rank = torch.distributed.get_rank()
        # tensors_gather[rank] = tensor
        # output = torch.cat(tensors_gather, dim=0)
        # torch.distributed.all_reduce(output)
        # print(rank, output[0], force=True)
        torch.distributed.all_gather(tensors_gather, tensor.detach(), async_op=False)
        output = torch.cat(tensors_gather, dim=0)
        return output
    except AssertionError as e:
        print('|WARNING :', e)
        return tensor

@register_model('backsum_transformer')
class BacksumTransformerModel(BaseFairseqModel):
    """Train Transformer models for multiple language pairs simultaneously.

    Requires `--task multilingual_translation`.

    We inherit all arguments from TransformerModel and assume that all language
    pairs use a single Transformer architecture. In addition, we provide several
    options that are specific to the multilingual setting.

    Args:
        --share-encoder-embeddings: share encoder embeddings across all source languages
        --share-decoder-embeddings: share decoder embeddings across all target languages
        --share-encoders: share all encoder params (incl. embeddings) across all source languages
        --share-decoders: share all decoder params (incl. embeddings) across all target languages
    """

    def __init__(self, args, encoders, decoders, task):
        super().__init__()
        assert encoders.keys() == decoders.keys()
        self.keys = list(encoders.keys())
        for key in self.keys:
            assert isinstance(encoders[key], FairseqEncoder)
            assert isinstance(decoders[key], FairseqIncrementalDecoder)

        self.models = nn.ModuleDict({
            key: FairseqEncoderDecoderModel(encoders[key], decoders[key])
            for key in self.keys
        })
        assert args.share_encoders
        if (task.momentum_contrast_loss_ratio > 0.0 or task.moco_steps is not None
        or task.byol_ratio > 0.0 or task.byol_steps is not None
        or task.parallel_byol_ratio > 0.0 or task.parallel_byol_steps is not None
        or task.gold_gen_byol_ratio > 0.0 or task.gold_gen_steps is not None
        or task.decoder_byol_ratio > 0.0 or task.decoder_byol_steps is not None):
            self.encoder_extra_fc = nn.Sequential(nn.Linear(args.encoder_embed_dim, args.encoder_embed_dim * 4), nn.ReLU(), nn.Linear(args.encoder_embed_dim * 4, args.encoder_embed_dim))
            if not args.cross_byol:
                self.encoder_extra_fc2 = nn.Sequential(LayerNorm(args.encoder_embed_dim), nn.Linear(args.encoder_embed_dim, args.encoder_embed_dim * 4), nn.ReLU(), nn.Linear(args.encoder_embed_dim * 4, args.encoder_embed_dim))

            self.momentum_encoder = copy.deepcopy(encoders[self.keys[0]])
            self.momentum_encoder.requires_grad_(False)
            if task.decoder_byol_ratio > 0.0 or task.decoder_byol_steps is not None:
                self.momentum_decoder = copy.deepcopy(decoders[self.keys[0]])
                self.momentum_decoder.requires_grad_(False)
            self.momentum_encoder_extra_fc = copy.deepcopy(self.encoder_extra_fc)
            self.momentum_encoder_extra_fc.requires_grad_(False)
        if task.momentum_contrast_loss_ratio > 0.0 or task.moco_steps is not None:
            momentum_bank = torch.nn.functional.normalize( torch.randn([args.encoder_embed_dim, args.momentum_contrast_capcity], device='cuda'), dim=0 )
            if args.fp16:
                momentum_bank = momentum_bank.half()

            self.register_buffer('momentum_bank', momentum_bank)
            self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))

        if (task.byol_ratio > 0.0 or task.byol_steps is not None
        or task.parallel_byol_ratio or task.parallel_byol_steps is not None
        or task.gold_gen_byol_ratio or task.gold_gen_steps is not None
        or task.decoder_byol_ratio > 0.0 or task.decoder_byol_steps is not None):
            from fairseq.modules.multihead_attention import MultiheadAttention
            self.cross_attention = MultiheadAttention(
                args.encoder_embed_dim,
                args.encoder_attention_heads,
                kdim=args.encoder_embed_dim,
                vdim=args.encoder_embed_dim,
                dropout=args.attention_dropout,
                encoder_decoder_attention=True,
            )

        # lang embedding for cnndm and inshorts, 0 for cnndm, 1 for inshorts .
        if args.share_decoders and args.share_encoders:
            self.lang_embed_token = Embedding(len(args.langs), args.encoder_embed_dim, padding_idx=None)
            for k in self.models.keys():
                self.models[k].encoder.lang_embed_token = self.lang_embed_token
                self.models[k].decoder.lang_embed_token = self.lang_embed_token

    @torch.no_grad()
    def dequeue_and_enqueue(self, keys):
        keys = concat_all_gather(keys)

        batch_size = keys.shape[0]

        ptr = int(self.queue_ptr)
        momentum_contrast_capcity = self.momentum_bank.shape[1]
        # if ptr + batch_size > momentum_contrast_capcity:
        #     tail = momentum_contrast_capcity - ptr
        #     head = ptr + batch_size - momentum_contrast_capcity
        #     self.momentum_bank[:, ptr: ] = keys[:tail].T
        #     self.momentum_bank[:, :head] = keys[-head:].T
        # else:
        # # replace the keys at ptr (dequeue and enqueue)
        #     self.momentum_bank[:, ptr:ptr + batch_size] = keys.T
        self.momentum_bank.copy_(torch.cat([self.momentum_bank[:, batch_size:], keys.T], dim=1))
        ptr = (ptr + batch_size) % momentum_contrast_capcity # move pointer

        self.queue_ptr[0] = ptr

    def byol_cross_attention(self, q, k):
        x, attn = self.cross_attention(
            query=k,
            key=q,
            value=q
        )
        return x

    def forward(self, src_tokens, src_lengths, prev_output_tokens, src_lang, tgt_lang, src_lang_tokens=None, tgt_lang_tokens=None, **kwargs):
        key = '{}-{}'.format(src_lang, tgt_lang)
        encoder_out = self.models[key].encoder(src_tokens, src_lengths, lang_tokens=src_lang_tokens, **kwargs)
        decoder_out = self.models[key].decoder(
                prev_output_tokens, encoder_out, lang_tokens=tgt_lang_tokens, **kwargs,
            )
        decoder_out[1]['model_key'] = key
        return decoder_out

    def get_normalized_probs(self, net_output, log_probs, sample=None):
        key = net_output[1]['model_key']
        return self.models[key].get_normalized_probs(net_output, log_probs, sample)

    @staticmethod
    def add_args(parser):
        """Add model-specific arguments to the parser."""
        TransformerModel.add_args(parser)
        parser.add_argument('--share-encoder-embeddings', action='store_true',
                            help='share encoder embeddings across languages')
        parser.add_argument('--share-decoder-embeddings', action='store_true',
                            help='share decoder embeddings across languages')
        parser.add_argument('--share-encoders', action='store_true',
                            help='share encoders across languages')
        parser.add_argument('--share-decoders', action='store_true',
                            help='share decoders across languages')
        parser.add_argument('--load-decoders', action='store_true')

    @classmethod
    def build_model(cls, args, task):
        """Build a new model instance."""

        # make sure all arguments are present in older models
        base_multilingual_architecture(args)

        if not hasattr(args, 'max_source_positions'):
            args.max_source_positions = 1024
        if not hasattr(args, 'max_target_positions'):
            args.max_target_positions = 1024

        # src_langs = [lang_pair.split('-')[0] for lang_pair in task.model_lang_pairs]
        # tgt_langs = [lang_pair.split('-')[1] for lang_pair in task.model_lang_pairs]
        src_langs = task.langs
        tgt_langs = task.langs

        if args.share_encoders:
            args.share_encoder_embeddings = True
        if args.share_decoders:
            args.share_decoder_embeddings = True

        def build_embedding(dictionary, embed_dim, path=None):
            num_embeddings = len(dictionary)
            padding_idx = dictionary.pad()
            emb = Embedding(num_embeddings, embed_dim, padding_idx)
            # if provided, load from preloaded dictionaries
            if path:
                embed_dict = utils.parse_embedding(path)
                utils.load_embedding(embed_dict, dictionary, emb)
            return emb

        # build shared embeddings (if applicable)
        shared_encoder_embed_tokens, shared_decoder_embed_tokens = None, None
        if args.share_all_embeddings:
            if args.encoder_embed_dim != args.decoder_embed_dim:
                raise ValueError(
                    '--share-all-embeddings requires --encoder-embed-dim to match --decoder-embed-dim')
            if args.decoder_embed_path and (
                    args.decoder_embed_path != args.encoder_embed_path):
                raise ValueError('--share-all-embeddings not compatible with --decoder-embed-path')
            shared_encoder_embed_tokens = build_embedding(
                task.dict,
                embed_dim=args.encoder_embed_dim,
                path=args.encoder_embed_path,
            )
            shared_decoder_embed_tokens = shared_encoder_embed_tokens
            args.share_decoder_input_output_embed = True
        else:
            if args.share_encoder_embeddings:
                shared_encoder_embed_tokens = (
                    build_embedding(
                        task.dict,
                        embed_dim=args.encoder_embed_dim,
                        path=args.encoder_embed_path,
                    )
                )
            if args.share_decoder_embeddings:
                shared_decoder_embed_tokens = (
                    build_embedding(
                        task.dict,
                        embed_dim=args.decoder_embed_dim,
                        path=args.decoder_embed_path,
                    )
                )

        # encoders/decoders for each language
        lang_encoders, lang_decoders = {}, {}

        def get_encoder(lang):
            if lang not in lang_encoders:
                if shared_encoder_embed_tokens is not None:
                    encoder_embed_tokens = shared_encoder_embed_tokens
                else:
                    encoder_embed_tokens = build_embedding(
                        task.dict, args.encoder_embed_dim, args.encoder_embed_path
                    )
                lang_encoders[lang] = TransformerEncoder(args, task.dict, encoder_embed_tokens)
            return lang_encoders[lang]

        def get_decoder(lang):
            if lang not in lang_decoders:
                if shared_decoder_embed_tokens is not None:
                    decoder_embed_tokens = shared_decoder_embed_tokens
                else:
                    decoder_embed_tokens = build_embedding(
                        task.dict, args.decoder_embed_dim, args.decoder_embed_path
                    )
                lang_decoders[lang] = TransformerDecoder(args, task.dict, decoder_embed_tokens)
            return lang_decoders[lang]

        # shared encoders/decoders (if applicable)
        shared_encoder, shared_decoder = None, None
        if args.share_encoders:
            shared_encoder = get_encoder(src_langs[0])
        if args.share_decoders:
            shared_decoder = get_decoder(tgt_langs[0])

        encoders, decoders = OrderedDict(), OrderedDict()
        for lang_pair in task.model_lang_pairs:
            src, tgt = lang_pair.split('-')
            encoders[lang_pair] = shared_encoder if shared_encoder is not None else get_encoder(src)
            decoders[lang_pair] = shared_decoder if shared_decoder is not None else get_decoder(tgt)

        return cls(args, encoders, decoders, task)

    def load_state_dict(self, state_dict, strict=True, args=None):
        models_key = [k for k in state_dict.keys() if k.startswith('models')]
        if len(models_key) != 0:
            state_dict_models = state_dict.copy()
            for k, _ in state_dict.items():
                if not k.startswith('models.'):
                    assert k.startswith('lang_embed_token.') or k.startswith('momentum') or k.startswith('encoder_extra_fc') or k == "queue_ptr" or k.startswith("cross_attention")
                    continue
                assert k.startswith('models.')
                lang_pair = k.split('.')[1]
                if lang_pair not in self.models:
                    del state_dict_models[k]
            # for k, _ in state_dict.items():
            #     for lang_pair in self.models:
            #         if k.startswith('models.{}.encoder'.format(lang_pair)):
            #             new_k = k.replace('models.{}.encoder'.format(lang_pair), 'momentum_encoders.{}'.format(lang_pair))
            #             if new_k not in state_dict:
            #                 state_dict_models[new_k] = state_dict[k]
            if "encoder_extra_fc2.3.weight" not in state_dict_models:
                for k in ["encoder_extra_fc2.{}.weight".format(i) for i in range(3)]:
                    if k in state_dict_models:
                        del state_dict_models[k]
                        del state_dict_models[k.replace('weight', "bias")]
            missing_keys, unexpected_keys = super().load_state_dict(state_dict_models, strict=False)
            # assert unexpected_keys == []
            print("----------------------- unexpected_keys --------------")
            print(unexpected_keys)
            print("----------------------missing_keys-------------------")
            print(missing_keys)
            for key in unexpected_keys:
                assert key.startswith('encoder_extra_fc2')
            for key in missing_keys:
                assert key.startswith('momentum') or key.startswith('encoder_extra_fc') or key == "queue_ptr" or key.startswith("cross_attention")
            if missing_keys != []:
                self.momentum_encoder.load_state_dict(self.models[self.keys[0]].encoder.state_dict())
        else:
            state_dict_subset = state_dict.copy()
            print('Load {} parameters'.format(len(state_dict_subset)))
            if not args.share_encoders:
                encoder_key = [k for k in state_dict_subset.keys() if k.startswith('encoder')]
                for key in encoder_key:
                    del state_dict_subset[key]
            if not args.share_decoders and not args.load_decoders:
                decoder_key = [k for k in state_dict_subset.keys() if k.startswith('decoder')]
                for key in decoder_key:
                    del state_dict_subset[key]
            print('Reserve {} parameters'.format(len(state_dict_subset)))
            if not "encoder.lang_embed_token.weight" in state_dict_subset and hasattr(self, 'lang_embed_token'):
                state_dict_subset["encoder.lang_embed_token.weight"] = self.lang_embed_token.weight
            if not "decoder.lang_embed_token.weight" in state_dict_subset and hasattr(self, 'lang_embed_token'):
                state_dict_subset["decoder.lang_embed_token.weight"] = self.lang_embed_token.weight
            for k, m in self.models.items():
                missing_keys, unexpected_keys = m.load_state_dict(state_dict_subset, False)
                if args.share_encoders and args.share_decoders:
                    assert missing_keys + unexpected_keys == []
                    break

            if hasattr(self, "momentum_encoder"):
                encoder_state = self.models[self.keys[0]].encoder.state_dict()
                encoder_state.pop('lang_embed_token.weight', None)
                self.momentum_encoder.load_state_dict(encoder_state)

class TransformerEncoder(FairseqEncoder):
    """
    Transformer encoder consisting of *args.encoder_layers* layers. Each layer
    is a :class:`TransformerEncoderLayer`.

    Args:
        args (argparse.Namespace): parsed command-line arguments
        dictionary (~fairseq.data.Dictionary): encoding dictionary
        embed_tokens (torch.nn.Embedding): input embedding
    """

    def __init__(self, args, dictionary, embed_tokens):
        super().__init__(dictionary)
        self.register_buffer('version', torch.Tensor([3]))

        self.dropout = args.dropout
        self.encoder_layerdrop = args.encoder_layerdrop

        embed_dim = embed_tokens.embedding_dim
        self.padding_idx = embed_tokens.padding_idx
        self.max_source_positions = args.max_source_positions

        self.embed_tokens = embed_tokens

        self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt(embed_dim)

        self.embed_positions = PositionalEmbedding(
            args.max_source_positions, embed_dim, self.padding_idx,
            learned=args.encoder_learned_pos,
        ) if not args.no_token_positional_embeddings else None

        self.layer_wise_attention = getattr(args, 'layer_wise_attention', False)

        self.layers = nn.ModuleList([])
        self.layers.extend([
            TransformerEncoderLayer(args)
            for i in range(args.encoder_layers)
        ])

        if args.encoder_normalize_before:
            self.layer_norm = LayerNorm(embed_dim)
        else:
            self.layer_norm = None
        if getattr(args, 'layernorm_embedding', False):
            self.layernorm_embedding = LayerNorm(embed_dim)
        else:
            self.layernorm_embedding = None

    def forward_embedding(self, src_tokens, lang_tokens):
        # embed tokens and positions
        x = embed = self.embed_scale * self.embed_tokens(src_tokens)
        if self.embed_positions is not None:
            x = embed + self.embed_positions(src_tokens)
        # add lang_embedding
        if lang_tokens is not None and hasattr(self, 'lang_embed_token'):
            x.add_(self.lang_embed_token(src_tokens.clone().fill_(lang_tokens)))
        if self.layernorm_embedding:
            x = self.layernorm_embedding(x)
        x = F.dropout(x, p=self.dropout, training=self.training)
        return x, embed

    def forward(self, src_tokens, src_lengths, cls_input=None, lang_tokens=None, return_all_hiddens=False, **unused):
        """
        Args:
            src_tokens (LongTensor): tokens in the source language of shape
                `(batch, src_len)`
            src_lengths (torch.LongTensor): lengths of each source sentence of
                shape `(batch)`
            return_all_hiddens (bool, optional): also return all of the
                intermediate hidden states (default: False).

        Returns:
            namedtuple:
                - **encoder_out** (Tensor): the last encoder layer's output of
                  shape `(src_len, batch, embed_dim)`
                - **encoder_padding_mask** (ByteTensor): the positions of
                  padding elements of shape `(batch, src_len)`
                - **encoder_embedding** (Tensor): the (scaled) embedding lookup
                  of shape `(batch, src_len, embed_dim)`
                - **encoder_states** (List[Tensor]): all intermediate
                  hidden states of shape `(src_len, batch, embed_dim)`.
                  Only populated if *return_all_hiddens* is True.
        """
        if self.layer_wise_attention:
            return_all_hiddens = True

        x, encoder_embedding = self.forward_embedding(src_tokens, lang_tokens)

        # B x T x C -> T x B x C
        x = x.transpose(0, 1)

        # compute padding mask
        encoder_padding_mask = src_tokens.eq(self.padding_idx)
        if not encoder_padding_mask.any():
            encoder_padding_mask = None

        encoder_states = [] if return_all_hiddens else None

        # encoder layers
        for layer in self.layers:
            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
            dropout_probability = random.uniform(0, 1)
            if not self.training or (dropout_probability > self.encoder_layerdrop):
                x = layer(x, encoder_padding_mask)
                if return_all_hiddens:
                    encoder_states.append(x)

        if self.layer_norm:
            x = self.layer_norm(x)
            if return_all_hiddens:
                encoder_states[-1] = x

        return EncoderOut(
            encoder_out=x,  # T x B x C
            encoder_padding_mask=encoder_padding_mask,  # B x T
            encoder_embedding=encoder_embedding,  # B x T x C
            encoder_states=encoder_states,  # List[T x B x C]
        )

    def reorder_encoder_out(self, encoder_out, new_order):
        """
        Reorder encoder output according to *new_order*.

        Args:
            encoder_out: output from the ``forward()`` method
            new_order (LongTensor): desired order

        Returns:
            *encoder_out* rearranged according to *new_order*
        """
        if encoder_out.encoder_out is not None:
            encoder_out = encoder_out._replace(
                encoder_out=encoder_out.encoder_out.index_select(1, new_order)
            )
        if encoder_out.encoder_padding_mask is not None:
            encoder_out = encoder_out._replace(
                encoder_padding_mask=encoder_out.encoder_padding_mask.index_select(0, new_order)
            )
        if encoder_out.encoder_embedding is not None:
            encoder_out = encoder_out._replace(
                encoder_embedding=encoder_out.encoder_embedding.index_select(0, new_order)
            )
        if encoder_out.encoder_states is not None:
            for idx, state in enumerate(encoder_out.encoder_states):
                encoder_out.encoder_states[idx] = state.index_select(1, new_order)
        return encoder_out

    def max_positions(self):
        """Maximum input length supported by the encoder."""
        if self.embed_positions is None:
            return self.max_source_positions
        return min(self.max_source_positions, self.embed_positions.max_positions())

    def buffered_future_mask(self, tensor):
        dim = tensor.size(0)
        if not hasattr(self, '_future_mask') or self._future_mask is None or self._future_mask.device != tensor.device:
            self._future_mask = torch.triu(utils.fill_with_neg_inf(tensor.new(dim, dim)), 1)
            if self._future_mask.size(0) < dim:
                self._future_mask = torch.triu(utils.fill_with_neg_inf(self._future_mask.resize_(dim, dim)), 1)
        return self._future_mask[:dim, :dim]

    def upgrade_state_dict_named(self, state_dict, name):
        """Upgrade a (possibly old) state dict for new versions of fairseq."""
        if isinstance(self.embed_positions, SinusoidalPositionalEmbedding):
            weights_key = '{}.embed_positions.weights'.format(name)
            if weights_key in state_dict:
                print('deleting {0}'.format(weights_key))
                del state_dict[weights_key]
            state_dict['{}.embed_positions._float_tensor'.format(name)] = torch.FloatTensor(1)
        for i in range(len(self.layers)):
            # update layer norms
            self.layers[i].upgrade_state_dict_named(state_dict, "{}.layers.{}".format(name, i))

        version_key = '{}.version'.format(name)
        if utils.item(state_dict.get(version_key, torch.Tensor([1]))[0]) < 2:
            # earlier checkpoints did not normalize after the stack of layers
            self.layer_norm = None
            self.normalize = False
            state_dict[version_key] = torch.Tensor([1])
        return state_dict

class TransformerDecoder(FairseqIncrementalDecoder):
    """
    Transformer decoder consisting of *args.decoder_layers* layers. Each layer
    is a :class:`TransformerDecoderLayer`.

    Args:
        args (argparse.Namespace): parsed command-line arguments
        dictionary (~fairseq.data.Dictionary): decoding dictionary
        embed_tokens (torch.nn.Embedding): output embedding
        no_encoder_attn (bool, optional): whether to attend to encoder outputs
            (default: False).
    """

    def __init__(self, args, dictionary, embed_tokens, no_encoder_attn=False):
        super().__init__(dictionary)
        self.register_buffer('version', torch.Tensor([3]))

        self.dropout = args.dropout
        self.decoder_layerdrop = args.decoder_layerdrop
        self.share_input_output_embed = args.share_decoder_input_output_embed

        input_embed_dim = embed_tokens.embedding_dim
        embed_dim = args.decoder_embed_dim
        self.output_embed_dim = args.decoder_output_dim

        self.padding_idx = embed_tokens.padding_idx
        self.max_target_positions = args.max_target_positions

        self.embed_tokens = embed_tokens

        self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt(embed_dim)

        self.project_in_dim = Linear(input_embed_dim, embed_dim, bias=False) if embed_dim != input_embed_dim else None

        self.embed_positions = PositionalEmbedding(
            args.max_target_positions, embed_dim, self.padding_idx,
            learned=args.decoder_learned_pos,
        ) if not args.no_token_positional_embeddings else None

        self.cross_self_attention = getattr(args, 'cross_self_attention', False)
        self.layer_wise_attention = getattr(args, 'layer_wise_attention', False)

        self.layers = nn.ModuleList([])
        self.layers.extend([
            TransformerDecoderLayer(args, no_encoder_attn)
            for _ in range(args.decoder_layers)
        ])

        self.adaptive_softmax = None

        self.project_out_dim = Linear(embed_dim, self.output_embed_dim, bias=False) \
            if embed_dim != self.output_embed_dim and not args.tie_adaptive_weights else None

        if args.adaptive_softmax_cutoff is not None:
            self.adaptive_softmax = AdaptiveSoftmax(
                len(dictionary),
                self.output_embed_dim,
                options.eval_str_list(args.adaptive_softmax_cutoff, type=int),
                dropout=args.adaptive_softmax_dropout,
                adaptive_inputs=embed_tokens if args.tie_adaptive_weights else None,
                factor=args.adaptive_softmax_factor,
                tie_proj=args.tie_adaptive_proj,
            )
        elif not self.share_input_output_embed:
            self.embed_out = nn.Parameter(torch.Tensor(len(dictionary), self.output_embed_dim))
            nn.init.normal_(self.embed_out, mean=0, std=self.output_embed_dim ** -0.5)

        if args.decoder_normalize_before and not getattr(args, 'no_decoder_final_norm', False):
            self.layer_norm = LayerNorm(embed_dim)
        else:
            self.layer_norm = None
        if getattr(args, 'layernorm_embedding', False):
            self.layernorm_embedding = LayerNorm(embed_dim)
        else:
            self.layernorm_embedding = None

    def forward(
        self,
        prev_output_tokens,
        encoder_out=None,
        incremental_state=None,
        features_only=False,
        lang_tokens=None,
        **extra_args
    ):
        """
        Args:
            prev_output_tokens (LongTensor): previous decoder outputs of shape
                `(batch, tgt_len)`, for teacher forcing
            encoder_out (optional): output from the encoder, used for
                encoder-side attention
            incremental_state (dict): dictionary used for storing state during
                :ref:`Incremental decoding`
            features_only (bool, optional): only return features without
                applying output layer (default: False).

        Returns:
            tuple:
                - the decoder's output of shape `(batch, tgt_len, vocab)`
                - a dictionary with any model-specific outputs
        """
        x, extra = self.extract_features(
            prev_output_tokens,
            lang_tokens=lang_tokens,
            encoder_out=encoder_out,
            incremental_state=incremental_state,
            **extra_args
        )
        if not features_only:
            x = self.output_layer(x)
        return x, extra

    def extract_features(
        self,
        prev_output_tokens,
        encoder_out=None,
        incremental_state=None,
        full_context_alignment=False,
        alignment_layer=None,
        alignment_heads=None,
        lang_tokens=None,
        **unused,
    ):
        """
        Similar to *forward* but only return features.

        Includes several features from "Jointly Learning to Align and
        Translate with Transformer Models" (Garg et al., EMNLP 2019).

        Args:
            full_context_alignment (bool, optional): don't apply
                auto-regressive mask to self-attention (default: False).
            alignment_layer (int, optional): return mean alignment over
                heads at this layer (default: last layer).
            alignment_heads (int, optional): only average alignment over
                this many heads (default: all heads).

        Returns:
            tuple:
                - the decoder's features of shape `(batch, tgt_len, embed_dim)`
                - a dictionary with any model-specific outputs
        """
        if alignment_layer is None:
            alignment_layer = len(self.layers) - 1

        # embed positions
        positions = self.embed_positions(
            prev_output_tokens,
            incremental_state=incremental_state,
        ) if self.embed_positions is not None else None

        if incremental_state is not None:
            prev_output_tokens = prev_output_tokens[:, -1:].clone()
            if positions is not None:
                positions = positions[:, -1:].clone()

        # embed tokens and positions
        x = self.embed_scale * self.embed_tokens(prev_output_tokens)

        if self.project_in_dim is not None:
            x = self.project_in_dim(x)

        if positions is not None:
            x += positions

        # add lang embedding
        if lang_tokens is not None and hasattr(self, 'lang_embed_token'):
            x += self.lang_embed_token(prev_output_tokens.clone().fill_(lang_tokens))

        if self.layernorm_embedding:
            x = self.layernorm_embedding(x)

        x = F.dropout(x, p=self.dropout, training=self.training)

        # B x T x C -> T x B x C
        x = x.transpose(0, 1)

        self_attn_padding_mask = None
        if self.cross_self_attention or prev_output_tokens.eq(self.padding_idx).any():
            self_attn_padding_mask = prev_output_tokens.eq(self.padding_idx)

        # decoder layers
        attn = None
        inner_states = [x]
        for idx, layer in enumerate(self.layers):
            encoder_state = None
            if encoder_out is not None:
                if self.layer_wise_attention:
                    encoder_state = encoder_out.encoder_states[idx]
                else:
                    encoder_state = encoder_out.encoder_out

            if incremental_state is None and not full_context_alignment:
                self_attn_mask = self.buffered_future_mask(x)
            else:
                self_attn_mask = None

            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
            dropout_probability = random.uniform(0, 1)
            if not self.training or (dropout_probability > self.decoder_layerdrop):
                x, layer_attn = layer(
                    x,
                    encoder_state,
                    encoder_out.encoder_padding_mask if encoder_out is not None else None,
                    incremental_state,
                    self_attn_mask=self_attn_mask,
                    self_attn_padding_mask=self_attn_padding_mask,
                    need_attn=(idx == alignment_layer),
                    need_head_weights=(idx == alignment_layer),
                )
                inner_states.append(x)
                if layer_attn is not None and idx == alignment_layer:
                    attn = layer_attn.float()

        if attn is not None:
            if alignment_heads is not None:
                attn = attn[:alignment_heads]

            # average probabilities over heads
            attn = attn.mean(dim=0)

        if self.layer_norm:
            x = self.layer_norm(x)

        # T x B x C -> B x T x C
        x = x.transpose(0, 1)

        if self.project_out_dim is not None:
            x = self.project_out_dim(x)

        return x, {'attn': attn, 'inner_states': inner_states}

    def output_layer(self, features, **kwargs):
        """Project features to the vocabulary size."""
        if self.adaptive_softmax is None:
            # project back to size of vocabulary
            if self.share_input_output_embed:
                return F.linear(features, self.embed_tokens.weight)
            else:
                return F.linear(features, self.embed_out)
        else:
            return features

    def max_positions(self):
        """Maximum output length supported by the decoder."""
        if self.embed_positions is None:
            return self.max_target_positions
        return min(self.max_target_positions, self.embed_positions.max_positions())

    def buffered_future_mask(self, tensor):
        dim = tensor.size(0)
        if (
            not hasattr(self, '_future_mask')
            or self._future_mask is None
            or self._future_mask.device != tensor.device
            or self._future_mask.size(0) < dim
        ):
            self._future_mask = torch.triu(utils.fill_with_neg_inf(tensor.new(dim, dim)), 1)
        return self._future_mask[:dim, :dim]

    def upgrade_state_dict_named(self, state_dict, name):
        """Upgrade a (possibly old) state dict for new versions of fairseq."""
        if isinstance(self.embed_positions, SinusoidalPositionalEmbedding):
            weights_key = '{}.embed_positions.weights'.format(name)
            if weights_key in state_dict:
                del state_dict[weights_key]
            state_dict['{}.embed_positions._float_tensor'.format(name)] = torch.FloatTensor(1)

        for i in range(len(self.layers)):
            # update layer norms
            layer_norm_map = {
                '0': 'self_attn_layer_norm',
                '1': 'encoder_attn_layer_norm',
                '2': 'final_layer_norm'
            }
            for old, new in layer_norm_map.items():
                for m in ('weight', 'bias'):
                    k = '{}.layers.{}.layer_norms.{}.{}'.format(name, i, old, m)
                    if k in state_dict:
                        state_dict['{}.layers.{}.{}.{}'.format(name, i, new, m)] = state_dict[k]
                        del state_dict[k]

        version_key = '{}.version'.format(name)
        if utils.item(state_dict.get(version_key, torch.Tensor([1]))[0]) <= 2:
            # earlier checkpoints did not normalize after the stack of layers
            self.layer_norm = None
            self.normalize = False
            state_dict[version_key] = torch.Tensor([1])

        return state_dict

@register_model_architecture('backsum_transformer', 'backsum_transformer')
def base_multilingual_architecture(args):
    base_architecture(args)
    args.share_encoder_embeddings = getattr(args, 'share_encoder_embeddings', False)
    args.share_decoder_embeddings = getattr(args, 'share_decoder_embeddings', False)
    args.share_encoders = getattr(args, 'share_encoders', False)
    args.share_decoders = getattr(args, 'share_decoders', False)

@register_model_architecture('backsum_transformer', 'backsum_transformer_bart_large')
def multilingual_transformer_bart_large(args):
    args.encoder_embed_path = getattr(args, 'encoder_embed_path', None)
    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 1024)
    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 4*1024)
    args.encoder_layers = getattr(args, 'encoder_layers', 12)
    args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 16)
    args.encoder_normalize_before = getattr(args, 'encoder_normalize_before', False)
    args.encoder_learned_pos = getattr(args, 'encoder_learned_pos', True)
    args.decoder_embed_path = getattr(args, 'decoder_embed_path', None)
    args.decoder_embed_dim = getattr(args, 'decoder_embed_dim', args.encoder_embed_dim)
    args.decoder_ffn_embed_dim = getattr(args, 'decoder_ffn_embed_dim', args.encoder_ffn_embed_dim)
    args.decoder_layers = getattr(args, 'decoder_layers', 12)
    args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 16)
    args.decoder_normalize_before = getattr(args, 'decoder_normalize_before', False)
    args.decoder_learned_pos = getattr(args, 'decoder_learned_pos', True)
    args.attention_dropout = getattr(args, 'attention_dropout', 0.)
    args.relu_dropout = getattr(args, 'relu_dropout', 0.)
    args.dropout = getattr(args, 'dropout', 0.1)
    args.max_target_positions = getattr(args, 'max_target_positions', 1024)
    args.max_source_positions = getattr(args, 'max_source_positions', 1024)
    args.adaptive_softmax_cutoff = getattr(args, 'adaptive_softmax_cutoff', None)
    args.adaptive_softmax_dropout = getattr(args, 'adaptive_softmax_dropout', 0)
    args.share_decoder_input_output_embed = getattr(args, 'share_decoder_input_output_embed', True)
    args.share_all_embeddings = getattr(args, 'share_all_embeddings', True)

    args.decoder_output_dim = getattr(args, 'decoder_output_dim', args.decoder_embed_dim)
    args.decoder_input_dim = getattr(args, 'decoder_input_dim', args.decoder_embed_dim)

    args.no_scale_embedding = getattr(args, 'no_scale_embedding', True)
    args.layernorm_embedding = getattr(args, 'layernorm_embedding', True)

    args.activation_fn = getattr(args, 'activation_fn', 'gelu')
    args.pooler_activation_fn = getattr(args, 'pooler_activation_fn', 'tanh')
    args.pooler_dropout = getattr(args, 'pooler_dropout', 0.0)
    base_multilingual_architecture(args)


"""Utilities for runnning dice through the shciscf plugin."""
import glob
import os
import struct
from typing import Tuple

import h5py
import numpy as np
from pyscf import scf  # pylint: disable=import-error

try:
    from pyscf.shciscf import shci  # pylint: disable=import-error
except (ModuleNotFoundError, ImportError):
    print("issues importing pyscf.shciscf")
    raise ImportError

from ipie.qmc.afqmc import AFQMC
from ipie.trial_wavefunction.particle_hole import ParticleHole
from ipie.utils.from_pyscf import generate_hamiltonian

def get_perm(from_orb: list, to_orb: list, di: list, dj: list) -> int:
    """Determine sign of permutation needed to align two determinants."""
    nmove = 0
    perm = 0
    for o in from_orb:
        io = np.where(dj == o)[0]
        perm += io - nmove
        nmove += 1
    nmove = 0
    for o in to_orb:
        io = np.where(di == o)[0]
        perm += io - nmove
        nmove += 1
    return perm % 2 == 1

def read_dice_wavefunction(filename):
    print(f"Reading Dice wavefunction from {filename}")
    with open(filename, "rb") as f:
        data = f.read()
    _chr = 1
    _int = 4
    _dou = 8
    ndets_in_file = struct.unpack("<I", data[:4])[0]
    norbs = struct.unpack("<I", data[4:8])[0]
    wfn_data = data[8:]
    coeffs = []
    occs = []
    start = 0
    print(f"Number of determinants in dets.bin : {ndets_in_file}")
    print(f"Number of orbitals : {norbs}")
    for _ in range(ndets_in_file):
        coeff = struct.unpack("<d", wfn_data[start : start + _dou])[0]
        coeffs.append(coeff)
        start += _dou
        occ_i = wfn_data[start : start + norbs]
        occ_lists = decode_dice_det(str(occ_i)[2:-1])
        occs.append(occ_lists)
        start += norbs
    print("Finished reading wavefunction from file.")
    oa, ob = zip(*occs)
    return (
        np.array(coeffs, dtype=np.complex128),
        np.array(oa, dtype=np.int32),
        np.array(ob, dtype=np.int32),
    )

def convert_phase(coeff0, occa_ref, occb_ref, verbose=False):
    print("Converting phase to account for abab -> aabb")
    ndets = len(coeff0)
    coeffs = np.zeros_like(coeff0)
    for i in range(ndets):
        if verbose and i % max(1, (int(0.1 * ndets))) == 0 and i > 0:
            done = float(i) / ndets
            print(f"convert phase {i}. Percent: {done}")
        # doubles = list(set(occa_ref[i]) & set(occb_ref[i]))
        occa0 = np.array(occa_ref[i])
        occb0 = np.array(occb_ref[i])

        count = 0
        for ocb in occb0:
            passing_alpha = np.where(occa0 > ocb)[0]
            count += len(passing_alpha)

        phase = (-1) ** count
        coeffs[i] = coeff0[i] * phase
    ixs = np.argsort(np.abs(coeffs))[::-1]
    coeffs = coeffs[ixs]
    occa = np.array(occa_ref)[ixs]
    occb = np.array(occb_ref)[ixs]

    return coeffs, occa, occb

def decode_dice_det(occs):
    occ_a = []
    occ_b = []
    for i, occ in enumerate(occs):
        if occ == "2":
            occ_a.append(i)
            occ_b.append(i)
        elif occ == "a":
            occ_a.append(i)
        elif occ == "b":
            occ_b.append(i)
    return occ_a, occ_b

def read_opdm(filename: str) -> np.ndarray:
    """Read DICE RDM from filename."""
    with open(filename) as f:
        data = f.readlines()
        nmo = int(data[0])
        opdm = np.zeros((nmo, nmo))
        for l in data[1:]:
            _l = l.split()
            i, j, v = int(_l[0]), int(_l[1]), float(_l[2])
            opdm[i, j] = v
            opdm[j, i] = v

    return opdm

def get_noons(opdm: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    eig_val, eig_vec = np.linalg.eigh(opdm)
    # Want natural ordering (i.e. high to low) for noons
    return eig_val[::-1], eig_vec[:, ::-1]

def get_active_space_from_noons(
    noons: np.ndarray,
    nelec: tuple,
    thresh: float = 0.02,
    lower_bound_scaling_factor: int = 1,
) -> Tuple[int, Tuple[int, int]]:
    """Attempt to determine active space from SHCI noons."""
    frozen = noons >= 2.0 - thresh
    inactive = noons < thresh
    active_orbs = np.where(~frozen & ~inactive)[0]
    print(f"frozen occupancies: {noons[frozen]}")
    print(f"active occupancies: {noons[active_orbs]}")
    # find orbitals that likely correspond to electrons in active space
    # reference wavefunction.
    # lower_bound_scaling_factor accounts for picking singly occupied orbitals.
    occ_in_active_space = np.where(
        (noons <= 2.0 - thresh) & (noons > 1 - lower_bound_scaling_factor * thresh)
    )[0]
    nalpha, nbeta = nelec
    frozen_fermi_level = occ_in_active_space[0]
    nalpha_active = nalpha - frozen_fermi_level
    nbeta_active = nbeta - frozen_fermi_level
    num_occ_orbs = len(occ_in_active_space)
    expected_occ_orbs = max(nalpha_active, nbeta_active)
    err_string = f"Consider changing value of lower_bound_scaling_factor: {num_occ_orbs} vs {expected_occ_orbs}"
    assert num_occ_orbs == expected_occ_orbs, err_string
    return (len(active_orbs), (nalpha_active, nbeta_active))

def run_shci_coarse(
    mf: scf.RHF,
    cas: Tuple[int, Tuple[int, int]],
    num_proc: int,
    eps=1e-4,
    num_sweep=3,
    mpi_cmd="mpirun -np",
    num_dets=100000,
) -> shci.SHCI:
    """Run coarse SHCI calculation in presumably large active space.

    Typically we run this first and obtain the SHCI 1RDM, from which one can
    compute NOONS and hopefully select a useful active space.
    """
    assert eps < 1e-3
    assert num_sweep < 5
    sweep_epsilon = np.logspace(-3, np.log10(eps), num_sweep)
    sweep_iter = np.linspace(0, 5, num_sweep, dtype=int)
    nact, nelec = cas
    mc1 = shci.SHCISCF(mf, nact, nelec)
    print(f"Coarse SHCI: nact = {mc1.ncas}, ncore = {mc1.ncore}")
    path = os.getcwd()
    mc1.fcisolver.scratchDirectory = path
    mc1.fcisolver.runtimeDir = path
    mc1.chkfile = "mcscf.chk"
    mc1.fcisolver.sweep_iter = sweep_iter
    mc1.fcisolver.sweep_epsilon = sweep_epsilon
    mc1.fcisolver.nPTiter = 0
    mc1.max_iter = 1
    mc1.fcisolver.maxIter = 20
    mc1.max_cycle_macro = 1
    mc1.internal_rotation = True
    mc1.fcisolver.nPTiter = 0  # Turns off PT calculation, i.e. no PTRDM.
    mc1.fcisolver.mpiprefix = mpi_cmd + f" {num_proc}"
    shci.dryrun(mc1)
    shci.writeSHCIConfFile(mc1.fcisolver, mc1.nelecas, False)
    with open(path + "/" + mc1.fcisolver.configFile, "a") as f:
        f.write("noio\n")
        f.write(f"writebestdeterminants {num_dets}\n\n")
    shci.executeSHCI(mc1.fcisolver)
    mc1.e_tot = shci.readEnergy(mc1.fcisolver)
    print(f"Coarse SHCI total energy = {mc1.e_tot}")
    return mc1

def run_shciscf(
    mf: scf.RHF,
    cas: Tuple[int, Tuple[int, int]],
    nat_orbs: np.ndarray,
    num_proc: int,
    eps=1e-5,
    num_sweep=3,
    mpi_cmd="mpirun -np",
    num_dets=1000000,
):
    """Helper function to run SHCISCF and write determinants to a format useable
    by ipie.
    """
    assert eps < 1e-3
    assert num_sweep < 5
    scf.chkfile.dump_scf(mf.mol, "mcscf.chk", mf.e_tot, mf.mo_energy, mf.mo_coeff, mf.mo_occ)
    mc1 = shci.SHCISCF(mf, cas[0], cas[1])
    print(f"SHCISCF: nact = {mc1.ncas}, ncore = {mc1.ncore}")
    path = os.getcwd()
    sweep_epsilon = np.logspace(-3, np.log10(eps), num_sweep)
    sweep_iter = np.linspace(0, 5, num_sweep, dtype=int)
    mc1.fcisolver.scratchDirectory = path
    mc1.fcisolver.runtimeDir = path
    mc1.chkfile = mf.chkfile
    mc1.fcisolver.sweep_iter = sweep_iter
    mc1.fcisolver.sweep_epsilon = sweep_epsilon
    mc1.max_iter = 20
    mc1.fcisolver.maxIter = 20
    mc1.max_cycle_macro = 20
    mc1.internal_rotation = True
    mc1.fcisolver.nPTiter = 100
    mc1.fcisolver.mpiprefix = mpi_cmd + f" {num_proc}"
    mc1.chkfile = mc1._scf.chkfile
    mc1.kernel(nat_orbs)
    shci.dryrun(mc1)
    shci.writeSHCIConfFile(mc1.fcisolver, mc1.nelecas, False)
    with open(path + "/" + mc1.fcisolver.configFile, "a") as f:
        f.write("noio\n")
        f.write(f"writebestdeterminants {num_dets}\n\n")
    shci.executeSHCI(mc1.fcisolver)
    bkp_files = glob.glob("*.bkp")
    for bkp_file in bkp_files:
        os.remove(bkp_file)
    mc1.e_tot = shci.readEnergy(mc1.fcisolver)
    print(f"Coarse SHCI total energy = {mc1.e_tot}")
    return mc1

def build_trial_from_shciscf(
    mf: scf.RHF,
    nmo: int,
    nelec: Tuple[int, int],
    num_procs: int,
    noons_fname: str = "noons.h5",
    noons_thresh: float = 0.02,
    num_dets_for_trial: int = 1000000,
    convert_det_phase: bool = False,
) -> Tuple[ParticleHole, shci.SHCI]:
    """Wrapper function to generate trial using SHCI to select the active space.

    No guarantees about if this works well or is a sensible thing to do. Users
    should inspect noons etc.
    """
    print(f"Running coarse SHCI calculation in active space: ({nelec}, {nmo})")
    shci_inst = run_shci_coarse(mf, (nmo, nelec), num_procs)
    opdm = read_opdm(f"spatial1RDM.0.0.txt")
    print(f"Determining active space using 1RDM in")
    noons, nat_orbs = get_noons(opdm)
    print(f"Writing noons to {noons_fname}")
    with h5py.File(f"{noons_fname}", "w") as fh5:
        fh5["noons"] = noons
        fh5["nat_orb"] = nat_orbs
    cas = get_active_space_from_noons(
        noons,
        nelec,
        thresh=noons_thresh,
    )
    print("active space determined from coarse SHCI noons: ", cas)
    print(f"Natural orbital unitary shape: {nat_orbs.shape}")
    C = np.copy(mf.mo_coeff)
    # Rotate MOs by natural orbitals in active space selected.
    ncore = shci_inst.ncore
    nact = shci_inst.ncas
    print(f"Number of core orbitals in coarse SHCI: {ncore}")
    print(f"Number of active orbitals in coarse SHCI: {nact}")
    print(
        f"Coarse SCHI active space slice in original full set of orbitals: [{ncore}:{ncore+nact}]."
    )
    Cfrzn = C[:, :ncore].copy()
    Cactv = np.dot(C[:, ncore : ncore + nact], nat_orbs)
    Cvirt = C[:, ncore + nact :].copy()
    init_guess = np.hstack([Cfrzn, Cactv, Cvirt])
    shci_inst = run_shciscf(mf, cas, init_guess, num_procs, eps=1e-5)
    wfn = read_dice_wavefunction("dets.bin")
    if convert_det_phase:
        wfn = convert_phase(*wfn)
    nelec = mf.mol.nelec
    nmo = mf.mo_coeff.shape[-1]
    trial = ParticleHole(wfn, nelec, nmo, num_dets_for_trial=num_dets_for_trial)
    return trial, shci_inst

def build_driver_from_shciscf(
    mf: scf.RHF,
    cas: Tuple[int, Tuple[int, int]],
    num_procs: int,
    num_dets_for_trial: int = 1000000,
    num_walkers: int = 100,
    noons_thresh: float = 0.02,
    convert_det_phase: bool = False,
    chol_cut=1e-8,
    seed=7,
) -> "AFQMC":
    nact, nelec_cas = cas
    trial, shci_inst = build_trial_from_shciscf(
        mf,
        nact,
        nelec_cas,
        num_procs,
        num_dets_for_trial=num_dets_for_trial,
        noons_thresh=noons_thresh,
        convert_det_phase=convert_det_phase,
    )

    ham = generate_hamiltonian(
        mf.mol, shci_inst.mo_coeff, mf.get_hcore(), shci_inst.mo_coeff, chol_cut=chol_cut
    )

    afqmc = AFQMC.build(mf.mol.nelec, ham, trial)
    return afqmc, shci_inst


from django.utils import timezone

class LogUserLastSeenMiddleware:
    def get_ip_address(self, request):
        return request.META.get("HTTP_X_FORWARDED_FOR", None) or request.META.get(
            "REMOTE_ADDR"
        )

    def get_route(self, request):
        try:
            return request.resolver_match.url_name[:64]
        except (AttributeError, TypeError):
            return None

    def __init__(self, get_response):
        self.get_response = get_response

    def __call__(self, request):
        response = self.get_response(request)
        if hasattr(request, "user") and request.user.is_authenticated:
            route = self.get_route(request)
            if route:
                request.user.last_seen_ip = self.get_ip_address(request)
                request.user.last_seen_route = route
                request.user.last_seen_at = timezone.now()
                request.user.save(
                    update_fields=["last_seen_ip", "last_seen_route", "last_seen_at"]
                )
        return response


# import dialogflow_v2 as df
import dialogflow_v2beta1 as df
# import google.cloud.dialogflow_v2.types as df_types

from hazard_elicitation.settings import KEYS
from util.text.languages import LanguageConfig

class DialogFlowClient:
    LANGUAGE = LanguageConfig.LANGUAGE
    PROJECT = KEYS.get('dialogflow_project')
    CONTEXT = 'projects/{}/agent/sessions/{}/contexts/{}'
    KB = 'projects/{}/knowledgeBases/{}'
    """
    # Unused: audio input and output
    AUDIO_ENCODING = df.types.audio_config_pb2.AudioEncoding.AUDIO_ENCODING_LINEAR_16
    AUDIO_SAMPLE_RATE = 48000
    """

    @staticmethod
    def preprocess(external_contexts: list, session_id: str):
        """
        Updates existing contexts with external contexts.
        @param external_contexts:   List of dictionaries containing contexts.
        @param session_id:          The session of the conversation.
        @return: SessionClient, Session, Contexts
        """
        session_client = df.SessionsClient()
        session = session_client.session_path(DialogFlowClient.PROJECT, session_id)

        contexts = []
        for context in external_contexts:
            context_name = context.get('name', '')
            context_parameters = context.get('parameters', {})
            context_lifespan = context.get('lifespan', 1)
            if context_name and context_lifespan:
                parameters = df.types.Struct()
                parameters.update(context_parameters)
                contexts.append(df.types.Context(
                    name=DialogFlowClient.CONTEXT.format(DialogFlowClient.PROJECT, session_id, context_name),
                    lifespan_count=context_lifespan,
                    parameters=parameters))

        return session_client, session, contexts

    @staticmethod
    def detect_event(event: str, external_contexts: list, session_id: str):
        """
        Detects an intent from event input. An event is usually send by an application.
        @param event:               Name of the event
        @param external_contexts:   A list of dictionaries containing external contexts.
        @param session_id:          The session of the conversation.
        @return: DialogFlowAnalysisResult
        """
        session_client, session, contexts = DialogFlowClient.preprocess(external_contexts, session_id)

        query_params = df.types.QueryParameters(contexts=contexts)
        event_input = df.types.EventInput(name=event, language_code=DialogFlowClient.LANGUAGE)
        query_input = df.types.QueryInput(event=event_input)

        return session_client.detect_intent(
            session=session,
            query_input=query_input,
            query_params=query_params)

    @staticmethod
    def detect_intent(text: str, external_contexts: list, session_id: str):
        """
        Detects an intent from text input. Text is usually send by the user via a text input.
        @param text:                Text from the user.
        @param external_contexts:   A list of dictionaries containing external contexts.
        @param session_id:          The session of the conversation.
        @return: DialogFlowAnalysisResult
        """
        session_client, session, contexts = DialogFlowClient.preprocess(external_contexts, session_id)

        sentiment_config = df.types.SentimentAnalysisRequestConfig(analyze_query_text_sentiment=True)
        query_params = df.types.QueryParameters(
            knowledge_base_names=[DialogFlowClient.KB.format(DialogFlowClient.PROJECT, 'NzEwOTcxNzA2MzEwNjU2MDAw')],
            sentiment_analysis_request_config=sentiment_config,
            contexts=contexts)
        text_input = df.types.TextInput(text=text, language_code=DialogFlowClient.LANGUAGE)
        query_input = df.types.QueryInput(text=text_input)

        # Requires fulfillment text in DF UI
        # output_audio_config = df.types.OutputAudioConfig(audio_encoding=DialogFlowClient.AUDIO_ENCODING)

        return session_client.detect_intent(
            session=session,
            query_input=query_input,
            query_params=query_params)
        # output_audio_config=output_audio_config)

    """
    Unused: audio input and output
    @staticmethod
    def detect_intent_audio(audio: bytes, session: str):
        session_client = df.SessionsClient()
        session = session_client.session_path(DialogFlowClient.PROJECT, session)

        audio_config = df.types.InputAudioConfig(
            audio_encoding=DialogFlowClient.AUDIO_ENCODING,
            language_code=DialogFlowClient.LANGUAGE,
            sample_rate_hertz=DialogFlowClient.AUDIO_SAMPLE_RATE)
        query_input = df.types.QueryInput(audio_config=audio_config)

        sentiment_config = df.types.SentimentAnalysisRequestConfig(analyze_query_text_sentiment=True)
        query_params = df.types.QueryParameters(sentiment_analysis_request_config=sentiment_config)
        output_audio_config = df.types.OutputAudioConfig(audio_encoding=DialogFlowClient.AUDIO_ENCODING)

        return session_client.detect_intent(
            session=session,
            query_input=query_input,
            query_params=query_params,
            input_audio=audio,
            output_audio_config=output_audio_config
        )

    @staticmethod
    def detect_intent_audio_stream(audio: bytes, session: str):
        session_client = df.SessionsClient()
        session = session_client.session_path(DialogFlowClient.PROJECT, session)

        audio_config = df.types.InputAudioConfig(
            audio_encoding=DialogFlowClient.AUDIO_ENCODING,
            language_code=DialogFlowClient.LANGUAGE,
            sample_rate_hertz=DialogFlowClient.AUDIO_SAMPLE_RATE)

        query_input = df.types.QueryInput(audio_config=audio_config)
        requests = [
            df.types.StreamingDetectIntentRequest(session=session, query_input=query_input),
            # TODO parse smaller chunks
            df.types.StreamingDetectIntentRequest(input_audio=audio)
        ]

        return session_client.streaming_detect_intent(requests=requests)
    """


from tests.test_base import BaseTestClass
import unittest

class ErrorTestCase(BaseTestClass):

    def test_error(self):
        res = self.client.get('/path_not_existent')

        print(f'\n\n RESPONSE {res}')
        self.assertEqual(404, res.status_code)

if __name__ == '__main__':
    unittest.main()


import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.models as models
import torchvision.transforms.functional as TVF

from itertools import chain
from typing import Optional, Union
import math

from .utils import ( fourier, create_fourier_basis, smooth_min )

class PositionalEncoder(nn.Module):
  def __init__(
    self,
    input_dims: int = 3,
    max_freq: float = 6.,
    N: int = 64,
    log_sampling: bool = False,
  ):
    super().__init__()
    if log_sampling:
      bands = 2**torch.linspace(1, max_freq, steps=N, requires_grad=False, dtype=torch.float)
    else:
      bands = torch.linspace(1, 2**max_freq, steps=N, requires_grad=False, dtype=torch.float)
    self.bands = nn.Parameter(bands, requires_grad=False)
    self.input_dims = input_dims
  def output_dims(self): return self.input_dims * 2 * len(self.bands)
  def forward(self, x):
    assert(x.shape[-1] == self.input_dims)
    raw_freqs = torch.tensordot(x, self.bands, dims=0)
    raw_freqs = raw_freqs.reshape(x.shape[:-1] + (-1,))
    return torch.cat([ raw_freqs.sin(), raw_freqs.cos() ], dim=-1)

class FourierEncoder(nn.Module):
  def __init__(
    self,
    input_dims: int = 3,
    # TODO rename this num freqs to be more accurate.
    freqs: int = 128,
    sigma: int = 1 << 5,
    device="cpu",
  ):
    super().__init__()
    self.input_dims = input_dims
    self.freqs = freqs
    self.basis = create_fourier_basis(freqs, features=input_dims, freq=sigma, device=device)
    self.basis = nn.Parameter(self.basis, requires_grad=False)
    self.extra_scale = 1
  def output_dims(self): return self.freqs * 2
  def forward(self, x): return fourier(x, self.extra_scale * self.basis)
  def scale_freqs(self, amt: 1+1e-5, cap=2):
    self.extra_scale *= amt
    self.extra_scale = min(self.extra_scale, cap)

class LearnedFourierEncoder(nn.Module):
  def __init__(
    self,
    input_dims: int = 3,
    num_freqs: int = 16,
    sigma: int = 1 << 5,
    device="cpu",
  ):
    super().__init__()
    self.input_dims = input_dims
    self.n_freqs = num_freqs
    self.basis = create_fourier_basis(num_freqs, features=input_dims, freq=sigma, device=device)
    self.basis = nn.Parameter(self.basis, requires_grad=False)
    self.extra_scale = nn.Parameter(torch.tensor(1, requires_grad=True), requires_grad=True)
  def output_dims(self): return self.n_freqs * 2
  def forward(self, x): return fourier(x, self.extra_scale * self.basis)

# It seems a cheap approximation to SIREN works just as well? Not entirely sure.
class NNEncoder(nn.Module):
  def __init__(
    self,
    input_dims: int = 3,
    out: int = 32,
    device=None,
  ):
    super().__init__()
    self.fwd = nn.Linear(input_dims, out)
  def output_dims(self): return self.fwd.out_features
  def forward(self, x):
    assert(x.shape[-1] == self.fwd.in_features)
    return torch.sin(30 * self.fwd(x))

# how to initialize the MLP, otherwise defaults to torch
mlp_init_kinds = {
  None,
  "zero",
  "kaiming",
  "siren",
  "xavier",
}

class SkipConnMLP(nn.Module):
  "MLP with skip connections and fourier encoding"
  def __init__(
    self,

    num_layers = 5,
    hidden_size = 256,
    in_size=3, out=3,

    skip=3,
    activation = nn.LeakyReLU(inplace=True),
    latent_size=0,

    enc: Optional[Union[FourierEncoder, PositionalEncoder, NNEncoder]] = None,

    # Record the last layers activation
    last_layer_act = False,

    linear=nn.Linear,

    init=None,
  ):
    assert(init in mlp_init_kinds), "Must use init kind"
    super(SkipConnMLP, self).__init__()
    self.in_size = in_size
    map_size = 0

    self.enc = enc
    if enc is not None: map_size += enc.output_dims()

    self.dim_p = in_size + map_size + latent_size
    self.skip = skip
    self.latent_size = latent_size
    skip_size = hidden_size + self.dim_p

    # with skip
    hidden_layers = [
      linear(
        skip_size if (i % skip) == 0 and i != num_layers-1 else hidden_size, hidden_size,
      ) for i in range(num_layers)
    ]

    self.init =  nn.Linear(self.dim_p, hidden_size)
    self.layers = nn.ModuleList(hidden_layers)
    self.out = nn.Linear(hidden_size, out)
    weights = [
      self.init.weight, self.out.weight,
      *[l.weight for l in self.layers],
    ]
    biases = [
      self.init.bias, self.out.bias,
      *[l.bias for l in self.layers],
    ]
    if init is None:
      ...
    elif init == "zero":
      for t in weights: nn.init.zeros_(t)
      for t in biases: nn.init.zeros_(t)
    elif init == "xavier":
      for t in weights: nn.init.xavier_uniform_(t)
      for t in biases: nn.init.zeros_(t)
    elif init == "siren":
      for t in weights:
        fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(t)
        a = math.sqrt(6 / fan_in)
        nn.init._no_grad_uniform_(t, -a, a)
      for t in biases: nn.init.zeros_(t)
    elif init == "kaiming":
      for t in weights: nn.init.kaiming_normal_(t, mode='fan_out')
      for t in biases: nn.init.zeros_(t)

    self.activation = activation
    self.last_layer_act = last_layer_act

  def forward(self, p, latent: Optional[torch.Tensor]=None):
    batches = p.shape[:-1]
    init = p.reshape(-1, p.shape[-1])

    if self.enc is not None: init = torch.cat([init, self.enc(init)], dim=-1)
    if self.latent_size != 0:
      assert(latent is not None), "Did not pass latent vector when some was expected"
      init = torch.cat([init, latent.reshape(-1, self.latent_size)], dim=-1)
    else: assert((latent is None) or (latent.shape[-1] == 0)), "Passed latent vector when none was expected"

    x = self.init(init)
    for i, layer in enumerate(self.layers):
      if i != len(self.layers)-1 and (i % self.skip) == 0:
        x = torch.cat([x, init], dim=-1)
      x = layer(self.activation(x))
    if self.last_layer_act: setattr(self, "last_layer_out", x.reshape(batches + (-1,)))
    out_size = self.out.out_features
    return self.out(self.activation(x)).reshape(batches + (out_size,))
  # smoothness of this sample along a given dimension for the last axis of a tensor
  def l2_smoothness(self, sample, values=None, noise=1e-1, dim=-1):
    if values is None: values = self(sample)
    adjusted = sample + noise * torch.rand_like(sample)
    adjusted = self(adjusted)
    return (values-adjusted).square().mean()
  def zero_last_layer(self):
    nn.init.zeros_(self.out.weight)
    nn.init.zeros_(self.out.bias)
  def uniform_last_layer(self, a=1e-4):
    nn.init.uniform_(self.out.weight, -a, a)
    nn.init.uniform_(self.out.bias, -a, a)
  # add an additional method for capt
  def variance(self, shape=None):
    return torch.stack([l.var(shape) for l in self.layers], dim=0)

class RecurrentUnit(nn.Module):
  def __init__(
    self,
    in_size=3,
    state_size:int = 128,
    enc: Optional[Union[FourierEncoder, PositionalEncoder, NNEncoder]] = None,
  ):
    super().__init__()
    self.enc = enc
    total_in = enc.output_dims() + in_size
    self.r_i = nn.Linear(total_in, state_size)
    self.r_s = nn.Linear(state_size, state_size)

    self.z_i = nn.Linear(total_in, state_size)
    self.z_s = nn.Linear(state_size, state_size)

    self.n_i = nn.Linear(total_in, state_size)
    self.n_s = nn.Linear(state_size, state_size)
  def forward(self, x, state):
    x = torch.cat([x, self.enc(x)], dim=-1)

    reset = torch.sigmoid(self.r_i(x) + self.r_s(state))
    update = torch.sigmoid(self.z_i(x) + self.z_s(state))
    new = torch.tanh(self.n_i(x) + reset * self.n_s(state))
    new_state = (1-update) * new + update * state
    return new_state

class EncodedGRU(nn.Module):
  def __init__(
    self,
    encs: [Union[FourierEncoder, PositionalEncoder, NNEncoder]],

    state_size:int=128,

    in_size:int=3,
    out:int=3,
    latent_out:int=128,
  ):
    super(EncodedGRU, self).__init__()
    self.in_size = in_size

    self.state_size = state_size

    hidden_layers = [
      RecurrentUnit(in_size=3, state_size=state_size, enc=enc) for enc in encs
    ]
    self.layers = nn.ModuleList(hidden_layers)

    self.last = nn.Linear(state_size, out + latent_out)
    self.out = out

  def forward(self, x, state=None):
    batches = x.shape[:-1]
    init = x.reshape(-1, x.shape[-1])

    if state is None:
      state = torch.zeros(*init.shape[:-1], self.state_size, device=x.device)

    out = []
    for l in self.layers:
      state = l(init, state)
      out.append(state[..., :self.out])
    last = self.last(state)
    out.append(last[..., :self.out])
    return torch.cat(out, dim=-1).reshape(*batches, -1), \
      last[..., self.out:].reshape(*batches, -1)

class Upsampler(nn.Module):
  def __init__(
    self,

    in_size: int,
    out: int,

    kernel_size:int = 3,

    repeat:int = 6,
    in_features:int = 3,
    out_features:int = 3,
    feat_decay: float = 2,
    activation = nn.LeakyReLU(inplace=True),
  ):
    super().__init__()
    step_size = (out - in_size)//repeat
    self.sizes = list(range(in_size + step_size, out+step_size, step_size))
    self.sizes = self.sizes[:repeat]
    self.sizes[-1] = out
    assert(kernel_size % 2 == 1), "Must provide odd kernel upsampling"

    feat_sizes = [
      max(out_features, int(in_features // (feat_decay**i))) for i in range(repeat+1)
    ]

    self.base = nn.Conv2d(in_features, out_features, kernel_size, 1, (kernel_size-1)//2)

    self.convs = nn.ModuleList([
      nn.Sequential(
        nn.Conv2d(fs, nfs, kernel_size, 1, (kernel_size-1)//2),
        nn.Dropout2d(0.1, inplace=True),
        nn.LeakyReLU(inplace=True)
      )
      for fs, nfs in zip(feat_sizes, feat_sizes[1:])
      # Move from one feature size to the next
    ])

    self.combine = nn.ModuleList([
      nn.Conv2d(feat_sizes, out_features, kernel_size, 1, (kernel_size-1)//2)
      for feat_sizes in feat_sizes[1:]
    ])

    self.rgb_up_kind="bilinear"
    self.feat_up_kind="nearest"

  def forward(self, x):
    curr = x.permute(0,3,1,2)
    upscaled = self.base(curr) # (N, H_in, W_in, 3)

    for s, conv, combine in zip(self.sizes, self.convs, self.combine):
      resized_old=F.interpolate(upscaled,size=(s,s),mode=self.rgb_up_kind,align_corners=False)

      curr = conv(F.interpolate(curr, size=(s, s), mode=self.feat_up_kind))
      upscaled = resized_old + combine(curr)
    return upscaled.permute(0,2,3,1)

# An update operator from https://github.com/princeton-vl/RAFT/blob/master/core/update.py
# takes as input a vector, with some hidden state and returns iterative updates on it.
class UpdateOperator(nn.Module):
  def __init__(
    self,
    in_size: int = 3,
    out_size: Optional[int] = None,
    hidden_size: int = 32,
    iters: int = 3,
  ):
    super().__init__()
    self.hidden_size = hs = hidden_size
    self.convz = nn.Conv3d(hs+in_size, hs, kernel_size=3, padding=1)
    self.convr = nn.Conv3d(hs+in_size, hs, kernel_size=3, padding=1)
    self.convq = nn.Conv3d(hs+in_size, hs, kernel_size=3, padding=1)

    self.conv1 = nn.Conv3d(hs, hs, 3, padding=1)
    self.conv2 = nn.Conv3d(hs, in_size, 3, padding=1)
    self.relu = nn.LeakyReLU(inplace=True)

    assert(iters > 0), "Must have at least one iteration"
    self.iters = iters
    self.out_size = out_size or in_size

  def forward(self, x, h=None):
    x = x.permute(1,4,0,2,3)
    if h is None:
      shape = list(x.shape)
      shape[1] = self.hidden_size
      h = torch.zeros(shape, device=x.device, dtype=x.dtype)
    else: raise NotImplementedError("TODO Have to permute h")
    init_x = x
    # total change in x
    for i in range(self.iters):
      hx = torch.cat([x.detach(), h], dim=1)
      z = self.convz(hx).sigmoid()
      r = self.convr(hx).sigmoid()
      q = self.convq(torch.cat([r*h, x], dim=1)).sigmoid()

      h = (1-z) * h + z * q
      dx = self.conv2(self.relu(self.conv1(h)))
      x = x + dx
    # returns total delta, rather than the shifted input
    out = (x - init_x).permute(2,0,3,4,1)[..., :self.out_size]
    return out

class SpatialEncoder(nn.Module):
  # Encodes an image into a latent vector, for use in PixelNeRF
  def __init__(
    self,
    latent_size: int =64,
    num_layers: int = 4,
  ):
    super().__init__()
    self.latent = None
    self.model = torchvision.models.resnet34(pretrained=True).eval()
    self.latent_size = latent_size
    self.num_layers = num_layers
  # (B, C = 3, H, W) -> (B, L, H, W)
  def forward(self, img):
    img = img.permute(0,3,1,2)
    l_sz = img.shape[2:]
    x = self.model.conv1(img)
    x = self.model.bn1(x)
    x = self.model.relu(x)
    latents = [x]
    # TODO other latent layers?

    ls = [F.interpolate(l, l_sz, mode="bilinear", align_corners=True) for l in latents]
    # necessary to detach here because we don't gradients to resnet
    # TODO maybe want latent gradients? Have to experiment with it.
    self.latents = torch.cat(latents, dim=1).detach().requires_grad_(False)
    return self.latents
  def sample(self, uvs: torch.Tensor, mode="bilinear", align_corners=True):
    assert(self.latents is not None), "Expected latents to be assigned in encoder"
    latents = F.grid_sample(
      self.latents,
      uvs.unsqueeze(0),
      mode=mode,
      align_corners=align_corners,
    )
    return latents.permute(0,2,3,1)

class Discriminator(nn.Module):
  def __init__(
    self,
    in_size=128,
    in_channels=3,
    hidden_size=64,
  ):
    super().__init__()
    self.in_size=128,
    hs = hidden_size
    ins = in_size
    assert(ins % 32 == 0)
    self.main = nn.Sequential(
      # input is 3 x ins x ins
      nn.Conv2d(in_channels, hidden_size, 4, 2, 1, bias=False),
      nn.LeakyReLU(inplace=True),
      # state size. (hs) x ins/2 x ins/2
      nn.Conv2d(hs, hs * 2, 4, 2, 1, bias=False),
      nn.BatchNorm2d(hs * 2),
      nn.LeakyReLU(inplace=True),
      # state size. (hs*2) x ins/4 x ins/4
      nn.Conv2d(hs * 2, hs * 4, 4, 2, 1, bias=False),
      nn.BatchNorm2d(hs * 4),
      nn.LeakyReLU(inplace=True),
      # state size. (hs*4) x ins/8 x ins/8
      nn.Conv2d(hs * 4, hs * 8, 4, 2, 1, bias=False),
      nn.BatchNorm2d(hs * 8),
      nn.LeakyReLU(inplace=True),
      # state size. (hs*8) x ins/16 x ins/16
      nn.Conv2d(hs * 8, hs * 16, 4, 2, 1, bias=False),
    )
    self.last = nn.Linear(hs * ins * ins / 64, 1)
  def forward(self, x, ref):
    fake = self.single_image_loss(x)
    real = self.single_image_loss(ref)
    raise NotImplementedError()
  def single_image_loss(self, x):
    # x: Batch, H, W, 3
    x = TVF.resize(x.permute(0, 3, 1, 2), size=(self.in_size, self.in_size))

    # out: Batch, 1
    return self.last(self.main(x))

def gram_matrix(img):
  batch, c, w, h = img.size()
  # a=batch size(=1)
  # b=number of feature maps
  # (c,d)=dimensions of a f. map (N=c*d)

  features = img.reshape(batch * c, w * h)  # resise F_XL into \hat F_XL

  G = torch.mm(features, features.t())  # compute the gram product

  # we 'normalize' the values of the gram matrix
  # by dividing by the number of element in each feature maps.
  return G.div(batch * c * w * h)

# performs N class classification of some set of points along with feature vectors.
class PointNet(nn.Module):
  def __init__(
    self,
    feature_size:int=7,
    classes:int=2,
    enc=None,
    intermediate_size=128,
  ):
    super().__init__()
    feats=feature_size
    self.first = SkipConnMLP(
      in_size=feats, out=intermediate_size,
      enc=enc, init="xavier",
      num_layers=3, skip=999
    )
    self.second = SkipConnMLP(
      in_size=intermediate_size * 2, out=classes,
      init="xavier", num_layers=3, skip=999
    )
  def forward(self, pos, feats):
    # input has shape (batches, num_samples, C)
    first = self.first(torch.cat([pos, feats], dim=-1))
    # smooth min & smooth max over items
    first_pool = -torch.logsumexp(-first, dim=-2)
    second_pool = torch.logsumexp(first, dim=-2)
    pooled = torch.cat([first_pool, second_pool], dim=-1)
    return self.second(pooled)

# create a module to normalize input image so we can easily put it in a
# nn.Sequential
class Normalization(nn.Module):
    def __init__(self):
        super().__init__()
        norm_mean = torch.tensor([0.485, 0.456, 0.406])
        norm_std = torch.tensor([0.229, 0.224, 0.225])
        # .view the mean and std to make them [C x 1 x 1] so that they can
        # directly work with image Tensor of shape [B x C x H x W].
        # B is batch size. C is number of channels. H is height and W is width.
        self.mean = nn.Parameter(norm_mean.reshape(-1, 1, 1), requires_grad=False)
        self.std = nn.Parameter(norm_std.reshape(-1, 1, 1), requires_grad=False)
    def forward(self, img):
      return (img - self.mean) / self.std

class StyleLoss(nn.Module):
  def __init__(self, feats):
    super().__init__()
    self.target = nn.Parameter(gram_matrix(feats).detach(), requires_grad=False)
  def forward(self, x):
    self.loss = F.mse_loss(gram_matrix(x), self.target)
    return x

class ContentLoss(nn.Module):
  def __init__(self, feats):
    super().__init__()
    self.feats = nn.Parameter(feats.detach(), requires_grad=False)
  def forward(self, x):
    self.loss = F.mse_loss(self.feats, x)
    return x

class StyleTransfer(nn.Module):
  def __init__(self, style_img, content_img=None):
    super().__init__()
    style_layers = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
    content_layers = ['conv_4']
    self.norm = Normalization()
    self.cnn = models.vgg19(pretrained=True).features.eval()
    # An ordered model of all the style and content losses
    model = nn.Sequential(self.norm)

    # list of a bunch of style loss modules
    style_losses = []
    content_losses = []

    i = 0
    for layer in self.cnn.children():
      if isinstance(layer, nn.Conv2d):
        i += 1
        name = 'conv_{}'.format(i)
      elif isinstance(layer, nn.ReLU):
        name = 'relu_{}'.format(i)
        # The in-place version doesn't play very nicely with the ContentLoss
        # and StyleLoss we insert below. So we replace with out-of-place
        # ones here.
        layer = nn.ReLU(inplace=False)
      elif isinstance(layer, nn.MaxPool2d): name = 'pool_{}'.format(i)
      elif isinstance(layer, nn.BatchNorm2d): name = 'bn_{}'.format(i)
      else: raise RuntimeError('Unrecognized layer: {}'.format(layer.__class__.__name__))

      model.add_module(name, layer)

      if name in style_layers:
        target_feature = model(style_img).detach()
        style_loss = StyleLoss(target_feature)
        model.add_module(f"style_loss_{i}", style_loss)
        style_losses.append(style_loss)
      if name in content_layers and content_img is not None:
        target_feature = model(content_img).detach()
        content_loss = ContentLoss(target_feature)
        model.add_module(f"content_loss_{i}", content_loss)
        content_losses.append(content_loss)
    # now we trim off the layers after the last content and style losses
    for i in range(len(model) - 1, -1, -1):
      if isinstance(model[i], StyleLoss): break
      if isinstance(model[i], ContentLoss): break

    self.model = model[:(i + 1)]
    self.style_losses = nn.ModuleList(style_losses)
    self.content_losses = nn.ModuleList(content_losses)

  def forward(self, x):
    self.model(x)
    sl_loss = 0
    cl_loss = 0
    for sl in self.style_losses: sl_loss += sl.loss
    for cl in self.content_losses: cl_loss += cl.loss
    return sl_loss, cl_loss

# https://arxiv.org/pdf/1903.09410.pdf
# Trying to parse hard to understand stuff.
class MonteCarloBNLinear(nn.Module):
  def __init__(self, in_features, out_features, bias=True, monte_carlo_samples:int=30):
    self.linear = nn.Linear(in_features, out_features,bias=bias);
    self.mc_samples = monte_carlo_samples
  def forward(self, x):
    if not self.training: return self.linear(x)
    x = x.expand(self.mc_samples, *x.shape)
    out = self.layers(x.reshape(-1, *x.shape[1:]))
    # training=True only checks that some parameters are correct but doesn't modify output
    out = F.batch_norm(out, torch.randn_like(out), torch.randn_like(out), training=True)
    out = out.reshape(self.mc_samples, *x.shape)
    self._var = out.stddev(dim=0)
    return out.mean(dim=0)
  def var(shape=None):
    if shape is None: return self.var
    return self.var.reshape_as(shape)


#!/usr/bin/env python3

import http.server
import io
import os
import re
import requests
import urllib.parse
from bs4 import BeautifulSoup
from pptx import Presentation
from pptx.dml.color import RGBColor
from pptx.util import Inches, Pt
from pptx.enum.text import MSO_VERTICAL_ANCHOR, MSO_AUTO_SIZE, PP_PARAGRAPH_ALIGNMENT

SHORT_TEXT_LIMIT_CHARS = 75
TITLE_FONT_PT = 75
SLIDE_BLANK_LAYOUT = 6
SLIDE_WIDTH_INCHES = 10
SLIDE_HEIGHT_INCHES = 7.5
SLIDE_SMALL_MARGIN_INCHES = 0.25
COLUMN_MARGIN_INCHES = 0.1
HEIGHT_MARGIN_INCHES = 0.1

# Init configuration default values
debug_logs = False
debug_slides = False
server_port = 8080

# Parse configuration environment variables
html2pptx_debug_logs = os.getenv('HTML2PPTX_DEBUG_LOGS', "false")
html2pptx_debug_slides = os.getenv('HTML2PPTX_DEBUG_SLIDES', "false")
html2pptx_server_port = os.getenv('HTML2PPTX_PORT', "8080")

# Set configuration variables depending on environment variables
if html2pptx_debug_logs.lower() == "true":
    debug_logs = True
if html2pptx_debug_slides.lower() == "true":
    debug_slides = True
server_port = int(html2pptx_server_port)

# Print configuration
print("Configuration:")
print("debug_logs:", debug_logs)
print("debug_slides:", debug_slides)
print("server_port:", server_port)

def html_to_pptx(url, css_selector):
    r = requests.get(url)
    url_string = r.text
    slides = html_to_slides(url_string, css_selector)
    prs_bytes_stream = slides_to_pptx(slides)
    return prs_bytes_stream

def html_to_slides(html_string, css_selector):
    soup = BeautifulSoup(html_string, 'html.parser')
    useful_content = soup.select(css_selector)
    slides = []
    for parent_content_tag in useful_content[0].children:
        if parent_content_tag.name is not None:
            slide_content = html_to_slide(parent_content_tag)
            slides.append(slide_content)
    return slides

def html_to_slide(parent_tag):
    return parse_tag_contents(parent_tag)

def parse_tag_contents(tag):
    tag_data = []
    for children_content_tag in tag.children:
        # Go through all children tags
        if children_content_tag.name is not None:
            # Just handle valid tags
            if children_content_tag.name == "img":
                # If we have an image, get the "src" link
                tag_data.append("img_src:" + children_content_tag["src"])
            elif children_content_tag.string is not None:
                # If we have only one string, return it
                if children_content_tag.string.strip() != "":
                    tag_data.append(children_content_tag.string.strip())
            else:
                # Get direct text elements from tag even if there are children elements with text inside
                # (but do not get the text from the children)
                direct_tag_strings = children_content_tag.find_all(string=True, recursive=False)
                sanitized_direct_tag_strings = []
                for string in direct_tag_strings:
                    sanitized_string = string.strip()
                    if re.match("^\\[if mso \\| IE\\].*", sanitized_string):
                        sanitized_string = ""
                    if sanitized_string != "":
                        sanitized_direct_tag_strings.append(sanitized_string)

                # Get direct text elements from tag even if there are children elements with text inside
                # (this time, we get the text from the children)
                recursive_tag_strings = children_content_tag.find_all(string=True, recursive=True)
                sanitized_recursive_tag_strings = []
                for string in recursive_tag_strings:
                    sanitized_string = string.strip()
                    if re.match("^\\[if mso \\| IE\\].*", sanitized_string):
                        sanitized_string = ""
                    if sanitized_string != "":
                        sanitized_recursive_tag_strings.append(sanitized_string)

                # If we have some direct text elements, then we are in a case of some formatted text nested within other
                # text tags, then just extract the whole text (direct and from children), return it,
                # and stop the recursion by going directly to the next element
                if len(sanitized_direct_tag_strings) > 0:
                    tag_data.append(" ".join(sanitized_recursive_tag_strings))
                    continue

                # If we are not in the case of nested text, just do a recursive call
                # to get the contents of the children tag
                tag_data.extend(parse_tag_contents(children_content_tag))
    return tag_data

def slides_to_pptx(slides):
    prs = Presentation()
    for slide in slides:
        if debug_logs:
            print("============================================ NEW SLIDE ============================================")
        fill_slide(prs, slide)
        if debug_logs:
            print("============================================ END SLIDE ============================================")
    prs_bytes_stream = io.BytesIO()
    prs.save(prs_bytes_stream)

    # Test to save bytes stream directly to file
    # Uncomment if needed
    # with open("test.pptx", 'wb') as out:
    #     out.write(prs_bytes_stream.getvalue())
    return prs_bytes_stream

def fill_slide(prs, slide):
    # Init default count values
    image_count = 0
    max_chars_in_strings = 0

    # Determine number of images and max string length
    for slide_data in slide:
        img_found = re.search("^img_src:(.*)$", slide_data)
        if img_found:
            image_count += 1
        else:
            if len(slide_data) > max_chars_in_strings:
                max_chars_in_strings = len(slide_data)

    # Determine if the slide is empty (no images and, max string length = 0, wich means no images and no text)
    empty_slide = image_count == 0 and max_chars_in_strings == 0
    if empty_slide:
        return

    # Add a slide when not empty
    prs_slide_layout = prs.slide_layouts[SLIDE_BLANK_LAYOUT]
    prs_slide = prs.slides.add_slide(prs_slide_layout)

    # Determine if we are in a column layout or not
    # We have a column layout if we have more than 1 image, at least 1 text block
    # and if the longest text is < SHORT_TEXT_LIMIT_CHARS chars
    with_multiple_images = image_count > 1
    with_short_texts = max_chars_in_strings != 0 and max_chars_in_strings <= SHORT_TEXT_LIMIT_CHARS
    column_layout = with_multiple_images and with_short_texts

    # Init base data
    images_array = []
    text_array = []
    column_text_array = []

    # Parse slide and separate images from text
    for slide_data in slide:
        image_found = re.search("^img_src:(.*)$", slide_data)
        if column_layout:
            # Special handling of text if we are in a column layout
            if image_found:
                # Be sure to always associate an image with some text below
                if len(images_array) == 0 and len(column_text_array) > 0:
                    # If text has been found without an image at the beginning, add an "empty" placeholder image
                    images_array.append("empty")
                    text_array.append(column_text_array)
                    column_text_array = []
                elif len(images_array) > 0:
                    # When some images have already been found, just add associated text
                    # to the previous image (empty or not)
                    text_array.append(column_text_array)
                    column_text_array = []
                images_array.append(image_found.group(1))
            else:
                # Add text to the current column
                column_text_array.append(slide_data)
        else:
            # Default handling when not in column layout
            if image_found:
                images_array.append(image_found.group(1))
            else:
                column_text_array.append(slide_data)

    if column_layout:
        # Always add the final text (empty or not) in column layout
        text_array.append(column_text_array)
    else:
        # Add text only if some text has been found (there will only be one big column in this case)
        if len(column_text_array) > 0:
            text_array.append(column_text_array)

    # Determine available space in slide and column layout
    available_slide_width_inches = SLIDE_WIDTH_INCHES - 2*SLIDE_SMALL_MARGIN_INCHES \
        - (len(images_array)-1)*COLUMN_MARGIN_INCHES
    column_width_inches = available_slide_width_inches
    if len(images_array) > 0:
        column_width_inches = available_slide_width_inches/len(images_array)

    # Init default values to compute image heights
    max_image_height_inches = 0
    images_heights_inches = []

    for index, image_link in enumerate(images_array):
        if debug_logs:
            print("IMAGE LINK:", image_link)

        if image_link == "empty":
            # If we have an "empty" placeholder image, ignore it but add a height of "0" to the image heights array
            images_heights_inches.append(0)
            continue

        # Determine image position
        top = Inches(SLIDE_SMALL_MARGIN_INCHES)
        left = Inches(SLIDE_SMALL_MARGIN_INCHES + index*column_width_inches
                      + index*COLUMN_MARGIN_INCHES)

        # Download image and add it to the slide
        image_req = requests.get(image_link)
        image_bytes = io.BytesIO(image_req.content)
        image_box = prs_slide.shapes.add_picture(image_bytes, left, top)

        # Determine image ratio and resize image (with aspect ratio preserved) to fit it in the slide if it is too large
        ratio = image_box.width.inches / image_box.height.inches
        if image_box.width.inches > SLIDE_WIDTH_INCHES - 2*SLIDE_SMALL_MARGIN_INCHES:
            image_box.width = Inches(SLIDE_WIDTH_INCHES - 2*SLIDE_SMALL_MARGIN_INCHES)
            image_box.height = Inches(image_box.width.inches / ratio)
        if image_box.height.inches > SLIDE_HEIGHT_INCHES - 2*SLIDE_SMALL_MARGIN_INCHES:
            image_box.height = Inches(SLIDE_HEIGHT_INCHES - 2*SLIDE_SMALL_MARGIN_INCHES)
            image_box.width = Inches(image_box.height.inches * ratio)

        # Center image horizontally if only one image is found
        if len(images_array) == 1:
            horizontal_image_center_inches = image_box.width.inches / 2
            slide_horizontal_center_inches = SLIDE_WIDTH_INCHES / 2
            left_horizontal_centered_inches = slide_horizontal_center_inches - horizontal_image_center_inches
            if left_horizontal_centered_inches < SLIDE_SMALL_MARGIN_INCHES:
                left_horizontal_centered_inches = SLIDE_SMALL_MARGIN_INCHES
            image_box.left = Inches(left_horizontal_centered_inches)

            # Center image vertically if this one image is alone with no text column
            if len(text_array) == 0:
                vertical_image_center_inches = image_box.height.inches / 2
                slide_vertical_center_inches = SLIDE_HEIGHT_INCHES / 2
                top_vertical_centered_inches = slide_vertical_center_inches - vertical_image_center_inches
                if top_vertical_centered_inches < SLIDE_SMALL_MARGIN_INCHES:
                    top_vertical_centered_inches = SLIDE_SMALL_MARGIN_INCHES
                image_box.top = Inches(top_vertical_centered_inches)

        images_heights_inches.append(image_box.height.inches)
        if image_box.height.inches > max_image_height_inches:
            max_image_height_inches = image_box.height.inches

    for index, text_column in enumerate(text_array):
        # For every text column, init default position and size
        left = Inches(SLIDE_SMALL_MARGIN_INCHES)
        width = Inches(SLIDE_WIDTH_INCHES - 2*SLIDE_SMALL_MARGIN_INCHES)
        top = Inches(SLIDE_SMALL_MARGIN_INCHES)
        height = Inches(SLIDE_HEIGHT_INCHES - 2*SLIDE_SMALL_MARGIN_INCHES)

        if len(images_array) > 0:
            # Override some default values if we have images
            top = Inches(SLIDE_SMALL_MARGIN_INCHES + max_image_height_inches + HEIGHT_MARGIN_INCHES)
            height = Inches(SLIDE_HEIGHT_INCHES - 2*SLIDE_SMALL_MARGIN_INCHES
                            - max_image_height_inches - HEIGHT_MARGIN_INCHES)

        if column_layout:
            # Column layout gets the final override if enabled
            left = Inches(SLIDE_SMALL_MARGIN_INCHES + index*column_width_inches
                          + index*COLUMN_MARGIN_INCHES)
            width = Inches(column_width_inches)
            top = Inches(SLIDE_SMALL_MARGIN_INCHES + images_heights_inches[index] + HEIGHT_MARGIN_INCHES)
            height = Inches(SLIDE_HEIGHT_INCHES - 2*SLIDE_SMALL_MARGIN_INCHES
                            - images_heights_inches[index] - HEIGHT_MARGIN_INCHES)

        # Create the text box
        text_box = prs_slide.shapes.add_textbox(left, top, width, height)

        if debug_slides:
            # Fill the shape with red in debug mode
            fill = text_box.fill
            fill.solid()
            fill.fore_color.rgb = RGBColor(255, 0, 0)

        # Create the text frame inside the text box and configure it
        text_frame = text_box.text_frame
        text_frame.clear()
        text_frame.vertical_anchor = MSO_VERTICAL_ANCHOR.MIDDLE
        if column_layout:
            # Change vertical anchor in column layout mode
            text_frame.vertical_anchor = MSO_VERTICAL_ANCHOR.TOP
        text_frame.word_wrap = True
        text_frame.auto_size = MSO_AUTO_SIZE.TEXT_TO_FIT_SHAPE

        # Add every string in the column to the text frame
        for text in text_column:
            # Determine if the text is a title (no images + only this text alone)
            is_title = len(images_array) == 0 and len(text_array) == 1 and len(text_column) == 1

            if debug_logs:
                print(text)

            # Fill the text frame
            paragraph = text_frame.paragraphs[0]
            if paragraph.text == "":
                if is_title:
                    # title format
                    paragraph.alignment = PP_PARAGRAPH_ALIGNMENT.CENTER
                    paragraph.font.size = Pt(TITLE_FONT_PT)
                paragraph.text = text
            else:
                paragraph = text_frame.add_paragraph()
                if is_title:
                    # title format
                    paragraph.alignment = PP_PARAGRAPH_ALIGNMENT.CENTER
                    paragraph.font.size = Pt(TITLE_FONT_PT)
                paragraph.text = text

class Html2pptx(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        # Handle GET requests

        # Set headers
        self.send_response(200)
        self.send_header("Content-type", "text/html")
        self.end_headers()

        # Send index.html page
        with open("index.html", 'rb') as out:
            self.wfile.write(out.read())

    def do_POST(self):
        # Handle POST requests

        # Retrieve and decode POST query data
        content_length = int(self.headers['Content-Length'])
        post_data = self.rfile.read(content_length)
        decoded_post_data = urllib.parse.parse_qs(post_data.decode("utf-8"))
        if debug_logs:
            print("decoded_post_data[\"url\"][0]", decoded_post_data["url"][0])
            print("decoded_post_data[\"selector\"][0]", decoded_post_data["selector"][0])

        # Translate HTML to PPTX, retrieves presentation bytes stream
        prs_bytes_stream = html_to_pptx(decoded_post_data["url"][0], decoded_post_data["selector"][0])

        # Set headers to download the PPTX file
        self.send_response(200)
        self.send_header("Content-Type", 'application/octet-stream')
        self.send_header("Content-Disposition", 'attachment; filename="presentation.pptx"')
        # This is some unused example code which uses content-length header when transferring a file
        # Since it seems to work here, we won't use it, but the code below will stay
        # here in case we need to use and modify it
        # Source:
        # https://stackoverflow.com/questions/18543640/how-would-i-create-a-python-web-server-that-downloads-a-file-on-any-get-request
        # fs = os.fstat(f.fileno())
        # self.send_header("Content-Length", str(fs.st_size))
        self.end_headers()

        # Send the PPTX presentation
        # Use getvalue() instead of read() with BytesIO to avoid problems
        # Source:
        # https://stackoverflow.com/questions/46981529/why-does-saving-a-presentation-to-a-file-like-object-produce-a-blank-presentatio?noredirect=1&lq=1
        self.wfile.write(prs_bytes_stream.getvalue())

# Setup and start HTTP server with custom Html2pptx handler
server_address = ("", server_port)
httpd = http.server.HTTPServer(server_address, Html2pptx)
print("Serving at port:", server_port)
httpd.serve_forever()