"""
Provide helper classes/functions to execute ELMo.
"""
# pylint: disable=no-self-use,arguments-differ,too-many-public-methods,too-many-lines
from typing import List, Tuple, Optional, Dict, Union, Any, Set
from collections import OrderedDict

import torch
from torch.nn.utils.rnn import PackedSequence
from torch.nn import ParameterList, Parameter

from pytorch_fast_elmo.factory import (
        ElmoCharacterEncoderFactory,
        ElmoWordEmbeddingFactory,
        ElmoLstmFactory,
        ElmoVocabProjectionFactory,
)
from pytorch_fast_elmo import utils

# TODO: use in inference.
# from _pytorch_fast_elmo import ScalarMix  # pylint: disable=no-name-in-module


def _raise_if_kwargs_is_invalid(allowed: Set[str], kwargs: Dict[str, Any]) -> None:
    invalid_keys = set(kwargs) - allowed
    if invalid_keys:
        msg = '\n'.join('invalid kwargs: {}'.format(key) for key in invalid_keys)
        raise ValueError(msg)


# Implement in Python as a temporary solution.
class ScalarMix(torch.nn.Module):  # type: ignore

    def __init__(
            self,
            mixture_size: int,
            do_layer_norm: bool = False,
            initial_scalar_parameters: Optional[List[float]] = None,
            trainable: bool = True,
    ) -> None:
        super().__init__()
        self.mixture_size = mixture_size
        self.do_layer_norm = do_layer_norm

        if initial_scalar_parameters is None:
            initial_scalar_parameters = [1.0 / mixture_size] * mixture_size
        elif len(initial_scalar_parameters) != mixture_size:
            raise ValueError("initial_scalar_parameters & mixture_size not match.")

        self.scalar_parameters = ParameterList([
                Parameter(
                        torch.FloatTensor([val]),
                        requires_grad=trainable,
                ) for val in initial_scalar_parameters
        ])
        self.gamma = Parameter(torch.FloatTensor([1.0]), requires_grad=trainable)

    def forward(
            self,
            tensors: List[torch.Tensor],  # pylint: disable=arguments-differ
            mask: torch.Tensor = None,
    ) -> torch.Tensor:

        def apply_layer_norm(tensor, broadcast_mask, num_elements_not_masked):  # type: ignore
            tensor_masked = tensor * broadcast_mask
            mean = torch.sum(tensor_masked) / num_elements_not_masked
            variance = torch.sum(torch.pow(
                    (tensor_masked - mean) * broadcast_mask,
                    2,
            )) / num_elements_not_masked
            return (tensor - mean) / torch.sqrt(variance + 1E-12)

        if len(tensors) != self.mixture_size:
            raise ValueError("tensors & mixture_size not match.")
        if self.do_layer_norm and mask is None:
            if tensors[0].ndimension() == 2:
                mask = torch.ones((tensors[0].shape[0],))
            else:
                raise ValueError("do_layer_norm but mask is not defined.")

        normed_weights = torch.split(
                torch.softmax(
                        torch.cat(list(self.scalar_parameters)),
                        0,
                ),
                1,
        )
        if self.do_layer_norm:
            mask_float = mask.float()
            broadcast_mask = mask_float.unsqueeze(-1)
            input_dim = tensors[0].size(-1)
            num_elements_not_masked = torch.sum(mask_float) * input_dim

        pieces = []
        for idx in range(self.mixture_size):
            tensor = tensors[idx]
            if self.do_layer_norm:
                tensor = apply_layer_norm(  # type: ignore
                        tensor,
                        broadcast_mask,
                        num_elements_not_masked,
                )
            weighted_tensor = normed_weights[idx] * tensor
            pieces.append(weighted_tensor)

        return self.gamma * sum(pieces)


class FastElmoBase(torch.nn.Module):  # type: ignore

    SCALAR_MIX_PARAMS = {
            'disable_scalar_mix',
            'num_output_representations',
            'output_representation_dropout',
            'scalar_mix_parameters',
            'do_layer_norm',
    }
    EXEC_PARAMS = {
            'exec_managed_lstm_bos_eos',
            'exec_managed_lstm_reset_states',
            'exec_sort_batch',
    }
    COMMON_PARAMS = SCALAR_MIX_PARAMS | EXEC_PARAMS

    _CHAR_CNN_FILTERS = [
            (1, 32),
            (2, 32),
            (3, 64),
            (4, 128),
            (5, 256),
            (6, 512),
            (7, 1024),
    ]

    def __init__(  # pylint: disable=dangerous-default-value
            self,

            # Generated by bilm-tf.
            options_file: Optional[str],
            weight_file: Optional[str],

            # Controls the behavior of execution.
            exec_managed_lstm_bos_eos: bool = True,
            exec_managed_lstm_reset_states: bool = False,
            exec_sort_batch: bool = True,

            # Controls the behavior of `ScalarMix`.
            disable_scalar_mix: bool = False,
            num_output_representations: int = 1,
            output_representation_dropout: float = 0.0,
            scalar_mix_parameters: Optional[List[float]] = None,
            do_layer_norm: bool = False,

            # Controls the behavior of factories.
            # Char CNN.
            disable_char_cnn: bool = False,
            char_cnn_requires_grad: bool = False,
            # From scratch.
            char_cnn_char_embedding_cnt: int = 261,
            char_cnn_char_embedding_dim: int = 16,
            char_cnn_filters: List[Tuple[int, int]] = _CHAR_CNN_FILTERS,
            char_cnn_activation: str = 'relu',
            char_cnn_num_highway_layers: int = 2,
            char_cnn_output_dim: int = 512,

            # Word Embedding.
            disable_word_embedding: bool = True,
            word_embedding_weight_file: Optional[str] = None,
            word_embedding_requires_grad: bool = False,
            # From scratch.
            word_embedding_cnt: int = 0,
            word_embedding_dim: int = 512,

            # The Forward LSTM.
            disable_forward_lstm: bool = False,
            forward_lstm_requires_grad: bool = False,
            # The Backward LSTM.
            disable_backward_lstm: bool = False,
            backward_lstm_requires_grad: bool = False,
            # From scratch.
            lstm_num_layers: int = 2,
            lstm_input_size: int = 512,
            lstm_hidden_size: int = 512,
            lstm_cell_size: int = 4096,
            lstm_cell_clip: float = 3.0,
            lstm_proj_clip: float = 3.0,
            lstm_truncated_bptt: int = 20,
            # Provide the BOS/EOS representations of shape `(projection_dim,)`
            # if char CNN is disabled.
            lstm_bos_repr: Optional[torch.Tensor] = None,
            lstm_eos_repr: Optional[torch.Tensor] = None,

            # The final softmax layer.
            disable_vocab_projection: bool = True,
            vocab_projection_requires_grad: bool = False,
            vocab_projection_input_size: int = 0,
            vocab_projection_proj_size: int = 0,
    ) -> None:
        super().__init__()

        self.disable_char_cnn = disable_char_cnn
        self.disable_word_embedding = disable_word_embedding
        self.disable_forward_lstm = disable_forward_lstm
        self.disable_backward_lstm = disable_backward_lstm
        self.disable_scalar_mix = disable_scalar_mix
        self.disable_vocab_projection = disable_vocab_projection

        self.exec_managed_lstm_bos_eos = exec_managed_lstm_bos_eos
        self.exec_managed_lstm_reset_states = exec_managed_lstm_reset_states
        self.exec_sort_batch = exec_sort_batch

        # Char CNN.
        if options_file:
            self.char_cnn_factory = ElmoCharacterEncoderFactory(
                    options_file,
                    weight_file,
            )
        else:
            # From scratch.
            self.char_cnn_factory = ElmoCharacterEncoderFactory.from_scratch(
                    char_embedding_cnt=char_cnn_char_embedding_cnt,
                    char_embedding_dim=char_cnn_char_embedding_dim,
                    filters=char_cnn_filters,
                    activation=char_cnn_activation,
                    num_highway_layers=char_cnn_num_highway_layers,
                    output_dim=char_cnn_output_dim,
            )

        if not disable_char_cnn:
            self._add_cpp_module_to_buffer(
                    'char_cnn',
                    self.char_cnn_factory.create(requires_grad=char_cnn_requires_grad),
            )

        # Word Embedding.
        if options_file:
            self.word_embedding_factory = ElmoWordEmbeddingFactory(
                    options_file,
                    word_embedding_weight_file or weight_file,
            )
        else:
            # From scratch.
            self.word_embedding_factory = ElmoWordEmbeddingFactory.from_scratch(
                    cnt=word_embedding_cnt,
                    dim=word_embedding_dim,
            )
            if exec_managed_lstm_bos_eos:
                raise ValueError('exec_managed_lstm_bos_eos should be disabled.')

        if not disable_word_embedding:
            # Not a cpp extension.
            word_embedding_weight, lstm_bos_repr, lstm_eos_repr = \
                    self.word_embedding_factory.create(requires_grad=word_embedding_requires_grad)
            self.register_buffer('word_embedding_weight', word_embedding_weight)

        # LSTM.
        if options_file:
            self.lstm_factory = ElmoLstmFactory(
                    options_file,
                    weight_file,
            )
        else:
            # From scratch.
            self.lstm_factory = ElmoLstmFactory.from_scratch(
                    num_layers=lstm_num_layers,
                    input_size=lstm_input_size,
                    hidden_size=lstm_hidden_size,
                    cell_size=lstm_cell_size,
                    cell_clip=lstm_cell_clip,
                    proj_clip=lstm_proj_clip,
                    truncated_bptt=lstm_truncated_bptt,
            )

        if not (disable_forward_lstm and disable_backward_lstm):
            forward_lstm, backward_lstm = self.lstm_factory.create(
                    enable_forward=not disable_forward_lstm,
                    forward_requires_grad=forward_lstm_requires_grad,
                    enable_backward=not disable_backward_lstm,
                    backward_requires_grad=backward_lstm_requires_grad,
            )
            if not disable_forward_lstm:
                self._add_cpp_module_to_buffer('forward_lstm', forward_lstm)
            if not disable_backward_lstm:
                self._add_cpp_module_to_buffer('backward_lstm', backward_lstm)

            # Cache BOS/EOS reprs.
            if exec_managed_lstm_bos_eos:
                if disable_char_cnn:
                    if lstm_bos_repr is None or lstm_eos_repr is None:
                        raise ValueError('BOS/EOS not provided.')

                else:
                    lstm_bos_repr, lstm_eos_repr = utils.get_bos_eos_token_repr(
                            self.char_cnn_factory,
                            self.char_cnn,
                    )
                self.register_buffer(
                        'lstm_bos_repr',
                        lstm_bos_repr,
                )
                self.register_buffer(
                        'lstm_eos_repr',
                        lstm_eos_repr,
                )

        # Vocabulary projection.
        if options_file:
            self.vocab_projection_factory = ElmoVocabProjectionFactory(
                    options_file,
                    weight_file,
            )
        else:
            self.vocab_projection_factory = ElmoVocabProjectionFactory.from_scratch(
                    vocab_projection_input_size,
                    vocab_projection_proj_size,
            )

        if not disable_vocab_projection:
            self.vocab_projection_weight, self.vocab_projection_bias = \
                    self.vocab_projection_factory.create(requires_grad=vocab_projection_requires_grad)

        # ScalarMix
        if not disable_scalar_mix:
            self.scalar_mixes: List[ScalarMix] = []

            for idx in range(num_output_representations):
                scalar_mix = ScalarMix(
                        # char cnn + lstm.
                        self.lstm_factory.num_layers + 1,
                        do_layer_norm=do_layer_norm,
                        initial_scalar_parameters=scalar_mix_parameters,
                        trainable=not scalar_mix_parameters,
                )
                self.add_module(f'scalar_mix_{idx}', scalar_mix)
                self.scalar_mixes.append(scalar_mix)

            self.repr_dropout = None
            if output_representation_dropout > 0.0:
                self.repr_dropout = torch.nn.Dropout(p=output_representation_dropout)

    def state_dict(  # type: ignore
            self,
            destination=None,
            prefix='',
            keep_vars=False,
    ):
        tmp_buffers = self._buffers
        self._buffers = OrderedDict()  # type: ignore
        ret = super().state_dict(destination, prefix, keep_vars)
        self._buffers = tmp_buffers
        return ret

    def _load_from_state_dict(  # type: ignore
            self,
            state_dict,
            prefix,
            local_metadata,
            strict,
            missing_keys,
            unexpected_keys,
            error_msgs,
    ):
        tmp_buffers = self._buffers
        self._buffers = OrderedDict()
        super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys,
                                      unexpected_keys, error_msgs)
        self._buffers = tmp_buffers

    def _add_cpp_module_to_buffer(self, name: str, cpp_module: Any) -> None:
        # register_buffer will raise an exception.
        self._buffers[name] = cpp_module

    def _get_lstm_device(self) -> int:
        cpp_ext = None
        if not self.disable_forward_lstm:
            cpp_ext = self.forward_lstm
        elif not self.disable_backward_lstm:
            cpp_ext = self.backward_lstm

        # Assume `cpp_ext` is not None.
        assert cpp_ext is not None
        tensor = cpp_ext.parameters()[0]
        return -1 if not tensor.is_cuda else tensor.get_device()  # type: ignore

    def get_batched_lstm_bos_eos_repr(self, attr_name: str, batch_size: int) -> PackedSequence:
        tensor = getattr(self, attr_name)

        if not tensor.is_cuda:
            # Move to GPU permanently.
            device = self._get_lstm_device()
            if device >= 0:
                tensor = tensor.cuda(device)
                setattr(self, attr_name, tensor)

        batched = tensor.unsqueeze(0).expand(batch_size, -1)
        return PackedSequence(batched, torch.LongTensor([batch_size]))

    def exec_forward_backword_lstm_bos_eos(
            self,
            lstm_attr_name: str,
            bos_eos_attr_name: str,
            batch_size: int,
    ) -> torch.Tensor:
        lstm = getattr(self, lstm_attr_name)
        batched = self.get_batched_lstm_bos_eos_repr(bos_eos_attr_name, batch_size)
        with torch.no_grad():
            outputs, _ = lstm(batched.data, batched.batch_sizes)
            # Returns the output of last layer.
            return outputs[-1]

    def exec_forward_lstm_bos(self, batch_size: int) -> torch.Tensor:
        return self.exec_forward_backword_lstm_bos_eos('forward_lstm', 'lstm_bos_repr', batch_size)

    def exec_forward_lstm_eos(self, batch_size: int) -> torch.Tensor:
        return self.exec_forward_backword_lstm_bos_eos('forward_lstm', 'lstm_eos_repr', batch_size)

    def exec_backward_lstm_bos(self, batch_size: int) -> torch.Tensor:
        return self.exec_forward_backword_lstm_bos_eos('backward_lstm', 'lstm_bos_repr', batch_size)

    def exec_backward_lstm_eos(self, batch_size: int) -> torch.Tensor:
        return self.exec_forward_backword_lstm_bos_eos('backward_lstm', 'lstm_eos_repr', batch_size)

    def exec_forward_lstm_permutate_states(self, index: torch.Tensor) -> None:
        self.forward_lstm.permutate_states(index)

    def exec_backward_lstm_permutate_states(self, index: torch.Tensor) -> None:
        self.backward_lstm.permutate_states(index)

    def exec_bilstm_permutate_states(self, index: torch.Tensor) -> None:
        if not self.disable_forward_lstm:
            self.exec_forward_lstm_permutate_states(index)
        if not self.disable_backward_lstm:
            self.exec_backward_lstm_permutate_states(index)

    def exec_char_cnn(self, inputs: PackedSequence) -> PackedSequence:
        """
        Char CNN.
        """
        output_data = self.char_cnn(inputs.data)
        return PackedSequence(output_data, inputs.batch_sizes)

    def exec_word_embedding(self, inputs: PackedSequence) -> PackedSequence:
        """
        Word embedding.
        """
        output_data = torch.nn.functional.embedding(
                inputs.data,
                self.word_embedding_weight,
                padding_idx=0,
        )
        return PackedSequence(output_data, inputs.batch_sizes)

    def exec_forward_lstm(
            self,
            inputs: PackedSequence,
    ) -> List[PackedSequence]:
        """
        Forward LSTM.
        """
        if self.exec_managed_lstm_bos_eos:
            max_batch_size = int(inputs.batch_sizes.data[0])
            # BOS.
            self.exec_forward_lstm_bos(max_batch_size)
        elif self.exec_managed_lstm_reset_states:
            self.forward_lstm.reset_states()

        # Feed inputs.
        outputs, _ = self.forward_lstm(inputs.data, inputs.batch_sizes)

        if self.exec_managed_lstm_bos_eos:
            # EOS.
            self.exec_forward_lstm_eos(max_batch_size)

        # To list of `PackedSequence`.
        return [PackedSequence(output, inputs.batch_sizes) for output in outputs]

    def exec_backward_lstm(
            self,
            inputs: PackedSequence,
    ) -> List[PackedSequence]:
        """
        Backward LSTM.
        """
        if self.exec_managed_lstm_bos_eos:
            max_batch_size = int(inputs.batch_sizes.data[0])
            # EOS.
            self.exec_backward_lstm_eos(max_batch_size)
        elif self.exec_managed_lstm_reset_states:
            self.backward_lstm.reset_states()

        # Feed inputs.
        outputs, _ = self.backward_lstm(inputs.data, inputs.batch_sizes)

        if self.exec_managed_lstm_bos_eos:
            # BOS.
            self.exec_backward_lstm_bos(max_batch_size)

        # To list of `PackedSequence`.
        return [PackedSequence(output, inputs.batch_sizes) for output in outputs]

    def exec_bilstm(
            self,
            inputs: PackedSequence,
    ) -> List[Tuple[PackedSequence, PackedSequence]]:
        """
        BiLSTM.
        """
        forward_seqs = self.exec_forward_lstm(inputs)
        backward_seqs = self.exec_backward_lstm(inputs)

        return list(zip(forward_seqs, backward_seqs))

    def concat_packed_sequences(
            self,
            packed_sequences: List[Tuple[PackedSequence, PackedSequence]],
    ) -> List[PackedSequence]:
        """
        Concatenate the outputs of fwd/bwd lstms.
        """
        return [
                PackedSequence(
                        torch.cat([fwd.data, bwd.data], dim=-1),
                        fwd.batch_sizes,
                ) for fwd, bwd in packed_sequences
        ]

    def combine_char_cnn_and_bilstm_outputs(
            self,
            char_cnn_packed: PackedSequence,
            bilstm_packed: List[PackedSequence],
    ) -> List[PackedSequence]:
        """
        Combine the outputs of Char CNN & BiLSTM for scalar mix.
        """
        # Simply duplicate the output of char cnn.
        duplicated_char_cnn_packed = PackedSequence(
                torch.cat([char_cnn_packed.data, char_cnn_packed.data], dim=-1),
                char_cnn_packed.batch_sizes,
        )

        combined = [duplicated_char_cnn_packed]
        combined.extend(bilstm_packed)
        return combined

    def exec_vocab_projection(self, context_repr: PackedSequence) -> PackedSequence:
        """
        Transform the last layer of LSTM to the probability distributions of vocabulary.
        """
        vocab_linear = torch.nn.functional.linear(
                context_repr.data,
                self.vocab_projection_weight,
                self.vocab_projection_bias,
        )
        vocab_probs = torch.nn.functional.softmax(vocab_linear, dim=-1)
        return PackedSequence(vocab_probs, context_repr.batch_sizes)

    def exec_scalar_mix(self, packed_sequences: List[PackedSequence]) -> List[PackedSequence]:
        """
        Scalar Mix.
        """
        reprs = []
        for scalar_mix in self.scalar_mixes:
            mixed = scalar_mix([inputs.data for inputs in packed_sequences])
            if self.repr_dropout is not None:
                mixed = self.repr_dropout(mixed)
            reprs.append(PackedSequence(mixed, packed_sequences[0].batch_sizes))
        return reprs

    def exec_bilstm_and_scalar_mix(
            self,
            token_repr: PackedSequence,
    ) -> List[PackedSequence]:
        """
        Common combination.
        """
        # BiLSTM.
        bilstm_repr = self.exec_bilstm(token_repr)
        # Scalar Mix.
        conbimed_repr = self.combine_char_cnn_and_bilstm_outputs(
                token_repr,
                self.concat_packed_sequences(bilstm_repr),
        )
        mixed_reprs = self.exec_scalar_mix(conbimed_repr)
        return mixed_reprs

    def pack_inputs(
            self,
            inputs: torch.Tensor,
            lengths: Optional[torch.Tensor] = None,
    ) -> PackedSequence:
        return utils.pack_inputs(inputs, lengths=lengths)

    def unpack_output(
            self,
            output: PackedSequence,
    ) -> torch.Tensor:
        return utils.unpack_outputs(output)

    def unpack_outputs(
            self,
            mixed_reprs: List[PackedSequence],
    ) -> List[torch.Tensor]:
        """
        Unpack the outputs of scalar mixtures.
        """
        return [self.unpack_output(mixed_repr) for mixed_repr in mixed_reprs]

    def to_allennlp_elmo_output_format(
            self,
            unpacks: List[torch.Tensor],
            mask: torch.Tensor,
    ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]:
        return {'elmo_representations': unpacks, 'mask': mask}

    def preprocess_inputs(
            self,
            inputs: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        lengths = utils.get_lengths_of_zero_padded_batch(inputs)
        original_lengths = lengths
        restoration_index: Optional[torch.Tensor] = None

        if self.exec_sort_batch:
            inputs, permutation_index, restoration_index = \
                    utils.sort_batch_by_length(inputs, lengths)
            lengths = lengths.index_select(0, permutation_index)
            self.exec_bilstm_permutate_states(permutation_index)

        return inputs, lengths, original_lengths, restoration_index

    def postprocess_outputs(
            self,
            unpacked_tensors: List[torch.Tensor],
            restoration_index: Optional[torch.Tensor],
            inputs: torch.Tensor,
            original_lengths: torch.Tensor,
    ) -> Tuple[List[torch.Tensor], torch.Tensor]:
        mask = utils.generate_mask_from_lengths(
                inputs.shape[0],
                inputs.shape[1],
                original_lengths,
        )
        if self.exec_sort_batch:
            assert restoration_index is not None
            unpacked_tensors = [
                    tensor.index_select(0, restoration_index) for tensor in unpacked_tensors
            ]
            self.exec_bilstm_permutate_states(restoration_index)

        return unpacked_tensors, mask

    def forward_with_sorting_and_packing(
            self,
            inputs: torch.Tensor,
    ) -> Tuple[List[torch.Tensor], torch.Tensor]:
        inputs, lengths, original_lengths, restoration_index = \
                self.preprocess_inputs(inputs)

        packed_inputs = self.pack_inputs(inputs, lengths)
        packed_outputs = self.execute(packed_inputs)

        unpacked_outputs = self.unpack_outputs(packed_outputs)
        unpacked_outputs, mask = self.postprocess_outputs(
                unpacked_outputs,
                restoration_index,
                inputs,
                original_lengths,
        )
        return unpacked_outputs, mask

    def execute(self, inputs: PackedSequence) -> List[PackedSequence]:
        raise NotImplementedError()

    def forward_like_allennlp(
            self,
            inputs: torch.Tensor,
    ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]:
        outputs, mask = self.forward_with_sorting_and_packing(inputs)
        return self.to_allennlp_elmo_output_format(outputs, mask)

    def forward(self):  # type: ignore
        raise NotImplementedError()


class FastElmo(FastElmoBase):

    def __init__(
            self,
            options_file: Optional[str],
            weight_file: str,
            **kwargs: Any,
    ) -> None:
        _raise_if_kwargs_is_invalid(
                self.COMMON_PARAMS | set([
                        # Fine-tuning is not fully supported by pytorch.
                        # 'char_cnn_requires_grad',
                        # 'forward_lstm_requires_grad',
                        # 'backward_lstm_requires_grad',
                ]),
                kwargs)
        super().__init__(options_file, weight_file, **kwargs)

    def execute(self, inputs: PackedSequence) -> List[PackedSequence]:
        token_repr = self.exec_char_cnn(inputs)
        mixed_reprs = self.exec_bilstm_and_scalar_mix(token_repr)
        return mixed_reprs

    def forward(  # type: ignore
            self,
            inputs: torch.Tensor,
    ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]:
        """
        The default workflow (same as AllenNLP).

        `inputs` of shape `(batch_size, max_timesteps, max_characters_per_token)
        """
        return self.forward_like_allennlp(inputs)


class FastElmoWordEmbedding(FastElmoBase):

    def __init__(
            self,
            options_file: Optional[str],
            weight_file: str,
            **kwargs: Any,
    ) -> None:
        _raise_if_kwargs_is_invalid(
                self.COMMON_PARAMS | {
                        'word_embedding_weight_file',
                        # Fine-tuning is not fully supported by pytorch.
                        # 'word_embedding_requires_grad',
                        # 'forward_lstm_requires_grad',
                        # 'backward_lstm_requires_grad',
                },
                kwargs)

        kwargs['disable_char_cnn'] = True
        kwargs['disable_word_embedding'] = False
        super().__init__(options_file, weight_file, **kwargs)

    def execute(self, inputs: PackedSequence) -> List[PackedSequence]:
        token_repr = self.exec_word_embedding(inputs)
        mixed_reprs = self.exec_bilstm_and_scalar_mix(token_repr)
        return mixed_reprs

    def forward(  # type: ignore
            self,
            inputs: torch.Tensor,
    ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]:
        """
        `inputs` of shape `(batch_size, max_timesteps)
        """
        return self.forward_like_allennlp(inputs)


class FastElmoPlainEncoderBase(FastElmoBase):  # pylint: disable=abstract-method

    def exec_context_independent_repr(self, inputs: PackedSequence) -> PackedSequence:
        raise NotImplementedError()

    def execute(self, inputs: PackedSequence) -> List[PackedSequence]:
        token_repr = self.exec_context_independent_repr(inputs)
        # BiLSTM.
        bilstm_repr = self.exec_bilstm(token_repr)
        # Scalar Mix.
        conbimed_repr = self.combine_char_cnn_and_bilstm_outputs(
                token_repr,
                self.concat_packed_sequences(bilstm_repr),
        )
        return conbimed_repr

    def forward(  # type: ignore
            self,
            inputs: torch.Tensor,
    ) -> Tuple[List[torch.Tensor], torch.Tensor]:
        """
        No scalar mix.
        """
        return self.forward_with_sorting_and_packing(inputs)


class FastElmoPlainEncoder(FastElmoPlainEncoderBase):

    def __init__(
            self,
            options_file: Optional[str],
            weight_file: str,
            **kwargs: Any,
    ) -> None:
        _raise_if_kwargs_is_invalid(self.EXEC_PARAMS, kwargs)
        kwargs['disable_scalar_mix'] = True
        super().__init__(options_file, weight_file, **kwargs)

    def exec_context_independent_repr(self, inputs: PackedSequence) -> PackedSequence:
        return self.exec_char_cnn(inputs)


class FastElmoWordEmbeddingPlainEncoder(FastElmoPlainEncoderBase):

    def __init__(
            self,
            options_file: Optional[str],
            weight_file: str,
            **kwargs: Any,
    ) -> None:
        _raise_if_kwargs_is_invalid(self.EXEC_PARAMS, kwargs)
        kwargs['disable_char_cnn'] = True
        kwargs['disable_word_embedding'] = False
        kwargs['disable_scalar_mix'] = True
        super().__init__(options_file, weight_file, **kwargs)

    def exec_context_independent_repr(self, inputs: PackedSequence) -> PackedSequence:
        return self.exec_word_embedding(inputs)


class FastElmoUnidirectionalVocabDistribBase(FastElmoBase):  # pylint: disable=abstract-method

    def exec_forward_vocab_prob_distrib(self, token_repr: PackedSequence) -> List[PackedSequence]:
        fwd_lstm_last = self.exec_forward_lstm(token_repr)[-1]
        fwd_vocab_distrib = self.exec_vocab_projection(fwd_lstm_last)
        return [fwd_vocab_distrib]

    def exec_backward_vocab_prob_distrib(self, token_repr: PackedSequence) -> List[PackedSequence]:
        bwd_lstm_last = self.exec_backward_lstm(token_repr)[-1]
        bwd_vocab_distrib = self.exec_vocab_projection(bwd_lstm_last)
        return [bwd_vocab_distrib]

    def forward(  # type: ignore
            self,
            inputs: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        (vocab_distrib,), mask = self.forward_with_sorting_and_packing(inputs)
        return vocab_distrib, mask


class FastElmoForwardVocabDistrib(FastElmoUnidirectionalVocabDistribBase):

    def __init__(
            self,
            options_file: Optional[str],
            weight_file: str,
            **kwargs: Any,
    ) -> None:
        _raise_if_kwargs_is_invalid(
                {'exec_managed_lstm_bos_eos'},
                kwargs,
        )
        kwargs['disable_backward_lstm'] = True
        kwargs['disable_vocab_projection'] = False
        kwargs['disable_scalar_mix'] = True
        super().__init__(options_file, weight_file, **kwargs)

    def execute(self, inputs: PackedSequence) -> List[PackedSequence]:
        token_repr = self.exec_char_cnn(inputs)
        return self.exec_forward_vocab_prob_distrib(token_repr)


class FastElmoBackwardVocabDistrib(FastElmoUnidirectionalVocabDistribBase):

    def __init__(
            self,
            options_file: Optional[str],
            weight_file: str,
            **kwargs: Any,
    ) -> None:
        _raise_if_kwargs_is_invalid(
                {'exec_managed_lstm_bos_eos'},
                kwargs,
        )
        kwargs['disable_forward_lstm'] = True
        kwargs['disable_vocab_projection'] = False
        kwargs['disable_scalar_mix'] = True
        super().__init__(options_file, weight_file, **kwargs)

    def execute(self, inputs: PackedSequence) -> List[PackedSequence]:
        token_repr = self.exec_char_cnn(inputs)
        return self.exec_backward_vocab_prob_distrib(token_repr)


class FastElmoWordEmbeddingForwardVocabDistrib(FastElmoUnidirectionalVocabDistribBase):

    def __init__(
            self,
            options_file: Optional[str],
            weight_file: str,
            **kwargs: Any,
    ) -> None:
        _raise_if_kwargs_is_invalid(
                {'exec_managed_lstm_bos_eos', 'word_embedding_weight_file'},
                kwargs,
        )
        kwargs['disable_char_cnn'] = True
        kwargs['disable_word_embedding'] = False
        kwargs['disable_backward_lstm'] = True
        kwargs['disable_vocab_projection'] = False
        kwargs['disable_scalar_mix'] = True
        super().__init__(options_file, weight_file, **kwargs)

    def execute(self, inputs: PackedSequence) -> List[PackedSequence]:
        token_repr = self.exec_word_embedding(inputs)
        return self.exec_forward_vocab_prob_distrib(token_repr)


class FastElmoWordEmbeddingBackwardVocabDistrib(FastElmoUnidirectionalVocabDistribBase):

    def __init__(
            self,
            options_file: Optional[str],
            weight_file: str,
            **kwargs: Any,
    ) -> None:
        _raise_if_kwargs_is_invalid(
                {'exec_managed_lstm_bos_eos', 'word_embedding_weight_file'},
                kwargs,
        )
        kwargs['disable_char_cnn'] = True
        kwargs['disable_word_embedding'] = False
        kwargs['disable_forward_lstm'] = True
        kwargs['disable_vocab_projection'] = False
        kwargs['disable_scalar_mix'] = True
        super().__init__(options_file, weight_file, **kwargs)

    def execute(self, inputs: PackedSequence) -> List[PackedSequence]:
        token_repr = self.exec_word_embedding(inputs)
        return self.exec_backward_vocab_prob_distrib(token_repr)
