import argparse
import copy
import os

from typing import Optional, Union

from ctranslate2.converters import utils
from ctranslate2.converters.converter import Converter
from ctranslate2.specs import common_spec, transformer_spec

_SUPPORTED_ACTIVATIONS = {
    "gelu": common_spec.Activation.GELUTanh,
    "relu": common_spec.Activation.RELU,
    "swish": common_spec.Activation.SWISH,
}


class OpenNMTTFConverter(Converter):
    """Converts OpenNMT-tf models."""

    @classmethod
    def from_config(
        cls,
        config: Union[str, dict],
        auto_config: bool = False,
        checkpoint_path: Optional[str] = None,
        model: Optional[str] = None,
    ):
        """Creates the converter from the configuration.

        Arguments:
          config: Path to the YAML configuration, or a dictionary with the loaded configuration.
          auto_config: Whether the model automatic configuration values should be used.
          checkpoint_path: Path to the checkpoint or checkpoint directory to load. If not set,
            the latest checkpoint from the model directory is loaded.
          model: If the model instance cannot be resolved from the model directory, this argument
             can be set to either the name of the model in the catalog or the path to the model
             configuration.

        Returns:
          A :class:`ctranslate2.converters.OpenNMTTFConverter` instance.
        """
        from opennmt import config as config_util
        from opennmt.utils.checkpoint import Checkpoint

        if isinstance(config, str):
            config = config_util.load_config([config])
        else:
            config = copy.deepcopy(config)

        if model is None:
            model = config_util.load_model(config["model_dir"])
        elif os.path.exists(model):
            model = config_util.load_model_from_file(model)
        else:
            model = config_util.load_model_from_catalog(model)

        if auto_config:
            config_util.merge_config(config, model.auto_config())

        data_config = config_util.try_prefix_paths(config["model_dir"], config["data"])
        model.initialize(data_config)

        checkpoint = Checkpoint.from_config(config, model)
        checkpoint_path = checkpoint.restore(checkpoint_path=checkpoint_path)
        if checkpoint_path is None:
            raise RuntimeError("No checkpoint was restored")

        model.create_variables()
        return cls(model)

    def __init__(self, model):
        """Initializes the converter.

        Arguments:
          model: An initialized and fully-built ``opennmt.models.Model`` instance.
        """
        self._model = model

    def _load(self):
        import opennmt

        if isinstance(self._model, opennmt.models.LanguageModel):
            spec_builder = TransformerDecoderSpecBuilder()
        else:
            spec_builder = TransformerSpecBuilder()

        return spec_builder(self._model)


class TransformerSpecBuilder:
    def __call__(self, model):
        import opennmt

        check = utils.ConfigurationChecker()
        check(
            isinstance(model, opennmt.models.Transformer),
            "Only Transformer models are supported",
        )
        check.validate()

        check(
            isinstance(model.encoder, opennmt.encoders.SelfAttentionEncoder),
            "Parallel encoders are not supported",
        )
        check(
            isinstance(
                model.features_inputter,
                (opennmt.inputters.WordEmbedder, opennmt.inputters.ParallelInputter),
            ),
            "Source inputter must be a WordEmbedder or a ParallelInputter",
        )
        check.validate()

        mha = model.encoder.layers[0].self_attention.layer
        ffn = model.encoder.layers[0].ffn.layer
        with_relative_position = mha.maximum_relative_position is not None
        activation_name = ffn.inner.activation.__name__

        check(
            activation_name in _SUPPORTED_ACTIVATIONS,
            "Activation %s is not supported (supported activations are: %s)"
            % (activation_name, ", ".join(_SUPPORTED_ACTIVATIONS.keys())),
        )
        check(
            with_relative_position != bool(model.encoder.position_encoder),
            "Relative position representation and position encoding cannot be both enabled "
            "or both disabled",
        )
        check(
            model.decoder.attention_reduction
            != opennmt.layers.MultiHeadAttentionReduction.AVERAGE_ALL_LAYERS,
            "Averaging all multi-head attention matrices is not supported",
        )

        source_inputters = _get_inputters(model.features_inputter)
        target_inputters = _get_inputters(model.labels_inputter)
        num_source_embeddings = len(source_inputters)
        if num_source_embeddings == 1:
            embeddings_merge = common_spec.EmbeddingsMerge.CONCAT
        else:
            reducer = model.features_inputter.reducer
            embeddings_merge = None
            if reducer is not None:
                if isinstance(reducer, opennmt.layers.ConcatReducer):
                    embeddings_merge = common_spec.EmbeddingsMerge.CONCAT
                elif isinstance(reducer, opennmt.layers.SumReducer):
                    embeddings_merge = common_spec.EmbeddingsMerge.ADD

            check(
                all(
                    isinstance(inputter, opennmt.inputters.WordEmbedder)
                    for inputter in source_inputters
                ),
                "All source inputters must WordEmbedders",
            )
            check(
                embeddings_merge is not None,
                "Unsupported embeddings reducer %s" % reducer,
            )

        alignment_layer = -1
        alignment_heads = 1
        if (
            model.decoder.attention_reduction
            == opennmt.layers.MultiHeadAttentionReduction.AVERAGE_LAST_LAYER
        ):
            alignment_heads = 0

        check.validate()

        encoder_spec = transformer_spec.TransformerEncoderSpec(
            len(model.encoder.layers),
            model.encoder.layers[0].self_attention.layer.num_heads,
            pre_norm=model.encoder.layer_norm is not None,
            activation=_SUPPORTED_ACTIVATIONS[activation_name],
            num_source_embeddings=num_source_embeddings,
            embeddings_merge=embeddings_merge,
            relative_position=with_relative_position,
        )

        decoder_spec = transformer_spec.TransformerDecoderSpec(
            len(model.decoder.layers),
            model.decoder.layers[0].self_attention.layer.num_heads,
            pre_norm=model.decoder.layer_norm is not None,
            activation=_SUPPORTED_ACTIVATIONS[activation_name],
            relative_position=with_relative_position,
            alignment_layer=alignment_layer,
            alignment_heads=alignment_heads,
        )

        spec = transformer_spec.TransformerSpec(encoder_spec, decoder_spec)

        spec.config.add_source_bos = bool(source_inputters[0].mark_start)
        spec.config.add_source_eos = bool(source_inputters[0].mark_end)
        for inputter in source_inputters:
            spec.register_source_vocabulary(_load_vocab(inputter.vocabulary_file))
        for inputter in target_inputters:
            spec.register_target_vocabulary(_load_vocab(inputter.vocabulary_file))

        self.set_transformer_encoder(
            spec.encoder,
            model.encoder,
            model.features_inputter,
        )
        self.set_transformer_decoder(
            spec.decoder,
            model.decoder,
            model.labels_inputter,
        )

        return spec

    def set_transformer_encoder(self, spec, module, inputter):
        for embedding_spec, inputter in zip(spec.embeddings, _get_inputters(inputter)):
            self.set_embeddings(embedding_spec, inputter)
        if module.position_encoder is not None:
            self.set_position_encodings(
                spec.position_encodings,
                module.position_encoder,
            )

        for layer_spec, layer in zip(spec.layer, module.layers):
            self.set_multi_head_attention(
                layer_spec.self_attention,
                layer.self_attention,
                self_attention=True,
            )

            self.set_ffn(layer_spec.ffn, layer.ffn)

        if module.layer_norm is not None:
            self.set_layer_norm(spec.layer_norm, module.layer_norm)

    def set_transformer_decoder(self, spec, module, inputter):
        self.set_embeddings(spec.embeddings, inputter)
        if module.position_encoder is not None:
            self.set_position_encodings(
                spec.position_encodings,
                module.position_encoder,
            )

        for layer_spec, layer in zip(spec.layer, module.layers):
            self.set_multi_head_attention(
                layer_spec.self_attention,
                layer.self_attention,
                self_attention=True,
            )

            if layer.attention:
                self.set_multi_head_attention(
                    layer_spec.attention,
                    layer.attention[0],
                    self_attention=False,
                )

            self.set_ffn(layer_spec.ffn, layer.ffn)

        if module.layer_norm is not None:
            self.set_layer_norm(spec.layer_norm, module.layer_norm)

        self.set_linear(spec.projection, module.output_layer)

    def set_ffn(self, spec, module):
        self.set_linear(spec.linear_0, module.layer.inner)
        self.set_linear(spec.linear_1, module.layer.outer)
        self.set_layer_norm_from_wrapper(spec.layer_norm, module)

    def set_multi_head_attention(self, spec, module, self_attention=False):
        split_layers = [common_spec.LinearSpec() for _ in range(3)]
        self.set_linear(split_layers[0], module.layer.linear_queries)
        self.set_linear(split_layers[1], module.layer.linear_keys)
        self.set_linear(split_layers[2], module.layer.linear_values)

        if self_attention:
            utils.fuse_linear(spec.linear[0], split_layers)
            if module.layer.maximum_relative_position is not None:
                spec.relative_position_keys = (
                    module.layer.relative_position_keys.numpy()
                )
                spec.relative_position_values = (
                    module.layer.relative_position_values.numpy()
                )
        else:
            utils.fuse_linear(spec.linear[0], split_layers[:1])
            utils.fuse_linear(spec.linear[1], split_layers[1:])

        self.set_linear(spec.linear[-1], module.layer.linear_output)
        self.set_layer_norm_from_wrapper(spec.layer_norm, module)

    def set_layer_norm_from_wrapper(self, spec, module):
        self.set_layer_norm(
            spec,
            module.output_layer_norm
            if module.input_layer_norm is None
            else module.input_layer_norm,
        )

    def set_layer_norm(self, spec, module):
        spec.gamma = module.gamma.numpy()
        spec.beta = module.beta.numpy()

    def set_linear(self, spec, module):
        spec.weight = module.kernel.numpy()
        if not module.transpose:
            spec.weight = spec.weight.transpose()
        if module.bias is not None:
            spec.bias = module.bias.numpy()

    def set_embeddings(self, spec, module):
        spec.weight = module.embedding.numpy()

    def set_position_encodings(self, spec, module):
        import opennmt

        if isinstance(module, opennmt.layers.PositionEmbedder):
            spec.encodings = module.embedding.numpy()[1:]


class TransformerDecoderSpecBuilder(TransformerSpecBuilder):
    def __call__(self, model):
        import opennmt

        check = utils.ConfigurationChecker()
        check(
            isinstance(model.decoder, opennmt.decoders.SelfAttentionDecoder),
            "Only self-attention decoders are supported",
        )
        check.validate()

        mha = model.decoder.layers[0].self_attention.layer
        ffn = model.decoder.layers[0].ffn.layer
        activation_name = ffn.inner.activation.__name__

        check(
            activation_name in _SUPPORTED_ACTIVATIONS,
            "Activation %s is not supported (supported activations are: %s)"
            % (activation_name, ", ".join(_SUPPORTED_ACTIVATIONS.keys())),
        )
        check.validate()

        spec = transformer_spec.TransformerDecoderModelSpec.from_config(
            len(model.decoder.layers),
            mha.num_heads,
            pre_norm=model.decoder.layer_norm is not None,
            activation=_SUPPORTED_ACTIVATIONS[activation_name],
        )

        spec.register_vocabulary(_load_vocab(model.features_inputter.vocabulary_file))
        self.set_transformer_decoder(
            spec.decoder,
            model.decoder,
            model.features_inputter,
        )

        return spec


def _get_inputters(inputter):
    import opennmt

    return (
        inputter.inputters
        if isinstance(inputter, opennmt.inputters.MultiInputter)
        else [inputter]
    )


def _load_vocab(vocab, unk_token="<unk>"):
    import opennmt

    if isinstance(vocab, opennmt.data.Vocab):
        tokens = list(vocab.words)
    elif isinstance(vocab, list):
        tokens = list(vocab)
    elif isinstance(vocab, str):
        tokens = opennmt.data.Vocab.from_file(vocab).words
    else:
        raise TypeError("Invalid vocabulary type")

    if unk_token not in tokens:
        tokens.append(unk_token)
    return tokens


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--config", help="Path to the YAML configuration.")
    parser.add_argument(
        "--auto_config",
        action="store_true",
        help="Use the model automatic configuration values.",
    )
    parser.add_argument(
        "--model_path",
        help=(
            "Path to the checkpoint or checkpoint directory to load. If not set, "
            "the latest checkpoint from the model directory is loaded."
        ),
    )
    parser.add_argument(
        "--model_type",
        help=(
            "If the model instance cannot be resolved from the model directory, "
            "this argument can be set to either the name of the model in the catalog "
            "or the path to the model configuration."
        ),
    )
    parser.add_argument(
        "--src_vocab",
        help="Path to the source vocabulary (required if no configuration is set).",
    )
    parser.add_argument(
        "--tgt_vocab",
        help="Path to the target vocabulary (required if no configuration is set).",
    )
    Converter.declare_arguments(parser)
    args = parser.parse_args()

    config = args.config
    if not config:
        if not args.model_path or not args.src_vocab or not args.tgt_vocab:
            raise ValueError(
                "Options --model_path, --src_vocab, --tgt_vocab are required "
                "when a configuration is not set"
            )

        model_dir = (
            args.model_path
            if os.path.isdir(args.model_path)
            else os.path.dirname(args.model_path)
        )
        config = {
            "model_dir": model_dir,
            "data": {
                "source_vocabulary": args.src_vocab,
                "target_vocabulary": args.tgt_vocab,
            },
        }

    converter = OpenNMTTFConverter.from_config(
        config,
        auto_config=args.auto_config,
        checkpoint_path=args.model_path,
        model=args.model_type,
    )
    converter.convert_from_args(args)


if __name__ == "__main__":
    main()