"""Specifications declare the expected variables layout of CTranslate2 models
that do not load a computation graph. The model converter should make sure that
each required variable of the specification is set.
"""

import abc
import ctypes
import json
import os
import shutil
import struct

from typing import Dict, List, Optional

import numpy as np

try:
    import torch

    torch_is_available = True
except ImportError:
    torch_is_available = False

OPTIONAL = "__optional"
CURRENT_BINARY_VERSION = 6

ACCEPTED_MODEL_TYPES = (
    "int8",
    "int8_float32",
    "int8_float16",
    "int8_bfloat16",
    "int16",
    "float16",
    "bfloat16",
    "float32",
)

SKIP_CREATING_ALIAS = ("rotary_scaling_long_factor", "rotary_scaling_short_factor")


def _join_scope(scope, name):
    if not scope:
        return name
    return "%s/%s" % (scope, name)


def _split_scope(scope):
    return scope.split("/")


def _parent_scope(scope):
    keys = _split_scope(scope)
    scope, attr = keys[:-1], keys[-1]
    return "/".join(scope), attr


def visit_spec(spec, fn, scope=""):
    """Recursively visits a layer spec."""
    for name, value in list(spec.__dict__.items()):
        if name.startswith("_"):
            continue
        if isinstance(value, list):
            for i, elem in enumerate(value):
                visit_spec(elem, fn, scope=_join_scope(scope, "%s_%d" % (name, i)))
        elif isinstance(value, LayerSpec):
            visit_spec(value, fn, scope=_join_scope(scope, name))
        else:
            fn(spec, _join_scope(scope, name), value)


def index_spec(spec, index):
    if not index:
        return spec
    keys = _split_scope(index)
    for key in keys:
        try:
            spec = getattr(spec, key)
        except AttributeError:
            attr, index = key.rsplit("_", 1)
            spec = getattr(spec, attr)[int(index)]
    return spec


class FrozenMeta(type):
    def __call__(self, *args, **kwargs):
        instance = super().__call__(*args, **kwargs)
        instance._frozen = True
        return instance


class FrozenAttr:
    def __setattr__(self, key, value):
        if hasattr(self, "_frozen") and not hasattr(self, key):
            raise AttributeError("Attribute %s does not exist" % key)
        super().__setattr__(key, value)


class LayerSpec(FrozenAttr, metaclass=FrozenMeta):
    """A layer specification declares the weights that should be set by the converters."""

    def validate(self) -> None:
        """Verify that the required weights are set.

        Raises:
          ValueError: If a required weight is not set in the specification.
        """
        unset_attributes = []

        def _check(spec, name, value):
            if value is None:
                unset_attributes.append(name)
                return

            if isinstance(value, np.ndarray):
                # float64 is not a supported type.
                if value.dtype == np.float64:
                    value = value.astype(np.float32)
            elif isinstance(value, float):
                value = np.dtype("float32").type(value)
            elif isinstance(value, bool):
                # Convert bool to an integer type.
                value = np.dtype("int8").type(value)
            elif isinstance(value, str):
                if value != OPTIONAL:
                    value = np.frombuffer(value.encode("utf-8"), dtype=np.int8)

            if isinstance(value, np.ndarray) or isinstance(value, np.generic):
                value = NumpyVariable(value)
            elif torch_is_available and isinstance(value, torch.Tensor):
                value = PyTorchVariable(value)

            attr_name = _split_scope(name)[-1]
            setattr(spec, attr_name, value)

        self._visit(_check)

        if unset_attributes:
            raise ValueError(
                "Some required model attributes are not set:\n\n%s"
                % "\n".join(unset_attributes)
            )

    def variables(
        self,
        prefix: str = "",
        ordered: bool = False,
    ) -> Dict[str, np.ndarray]:
        """Recursively returns the weights from this layer and its children.

        Arguments:
          prefix: Prefix to prepend to all variable names.
          ordered: If set, an ordered list is returned instead.

        Returns:
          Dictionary mapping variables name to value.
        """
        var = {}

        def _register_var(spec, name, value):
            if isinstance(value, str) and value == OPTIONAL:
                return
            var[_join_scope(prefix, name)] = value

        self._visit(_register_var)
        if ordered:
            return list(sorted(var.items(), key=lambda x: x[0]))
        return var

    def _alias_variables(self):
        """Find duplicate variables in spec and create aliases."""
        # When a variable is duplicated, keep the version that comes first in
        # the alphabetical order and alias the others.
        variables = self.variables(ordered=True)
        for name, value in reversed(variables):
            for other_name, other_value in variables:
                if name == other_name:
                    break
                # Because variables can be transformed on load (e.g. transposed),
                # we use an element-wise equality check.
                scope, attr_name = _parent_scope(name)
                if (
                    not value.is_scalar()
                    and value.equal(other_value)
                    and attr_name not in SKIP_CREATING_ALIAS
                ):
                    # Replace variable value by the alias name.
                    spec = index_spec(self, scope)
                    setattr(spec, attr_name, other_name)
                    break

    def _quantize(self, quantization):
        """Possibly quantizes the variable of the layer."""
        if quantization is not None and quantization not in ACCEPTED_MODEL_TYPES:
            raise ValueError(
                "%s is not a valid quantization type. Accepted types are: %s"
                % (quantization, ", ".join(ACCEPTED_MODEL_TYPES))
            )

        def _quantize(spec, name, value):
            if not isinstance(value, Variable) or value.is_scalar():
                return

            key = _split_scope(name)[-1]
            scale = None
            is_quantizable = hasattr(spec, "%s_scale" % key)
            is_convertible = value.dtype in ("float32", "float16", "bfloat16")

            if is_quantizable:
                if quantization == "int16":
                    value = value.to("float32").numpy()
                    # Represent the value with 10 bits so the multiplication is 20 bits
                    # and 12 bits are left for accumulation.
                    scale = np.float32(2**10 / np.amax(np.absolute(value)))
                    value *= scale
                    value = np.rint(value)
                    value = np.clip(
                        value, np.iinfo(np.int16).min, np.iinfo(np.int16).max
                    )
                    value = value.astype(np.int16)
                    scale = NumpyVariable(scale)
                    value = NumpyVariable(value)
                elif quantization in (
                    "int8",
                    "int8_float32",
                    "int8_float16",
                    "int8_bfloat16",
                ):
                    value = value.to("float32").numpy()
                    # For conv1d layer we need to reshape to 2D before calculating scale
                    old_shape = None
                    if len(value.shape) == 3:
                        old_shape = value.shape
                        value = value.reshape(value.shape[0], -1)
                    amax = np.amax(np.absolute(value), axis=1)
                    amax[amax == 0] = 127.0
                    scale = 127.0 / amax
                    value *= np.expand_dims(scale, 1)
                    value = np.rint(value)
                    value = value.astype(np.int8)
                    # reshape back to old shape
                    if old_shape:
                        value = value.reshape(old_shape)
                    scale = NumpyVariable(scale)
                    value = NumpyVariable(value)
                elif quantization in ("float16", "bfloat16", "float32"):
                    value = value.to(quantization)

            elif is_convertible:
                if quantization in ("float16", "int8_float16"):
                    value = value.to("float16")
                elif quantization in ("bfloat16", "int8_bfloat16"):
                    value = value.to("bfloat16")
                elif quantization in ("float32", "int16", "int8_float32"):
                    value = value.to("float32")

            setattr(spec, key, value)
            if scale is not None:
                setattr(spec, "%s_scale" % key, scale)

        self._visit(_quantize)

    def optimize(self, quantization: Optional[str] = None) -> None:
        """Recursively applies some optimizations to this layer:

        * Alias variables with the same shape and value.
        * Quantize weights.

        Arguments:
          quantization: Weight quantization scheme (possible values are: int8, int8_float32,
            int8_float16, int8_bfloat16, int16, float16, bfloat16, float32).
        """
        self._alias_variables()
        self._quantize(quantization)

    def _visit(self, fn):
        """Recursively visits this layer and its children."""
        visit_spec(self, fn)


def _dtype_to_type_id(object_dtype):
    # Order should match the DataType enum in include/ctranslate2/types.h
    dtypes = ("float32", "int8", "int16", "int32", "float16", "bfloat16")
    try:
        return dtypes.index(object_dtype)
    except ValueError:
        raise ValueError(
            "%s is not in list of supported dtypes: %s"
            % (object_dtype, ", ".join(dtypes))
        )


class ModelConfig(FrozenAttr, metaclass=FrozenMeta):
    """Base class for model configurations."""

    def __init__(self, **kwargs):
        """Initializes the configuration with a set of parameters."""
        for key, value in kwargs.items():
            setattr(self, key, value)

    def to_dict(self):
        """Returns the configuration as a dictionary."""
        return {
            key: value
            for key, value in self.__dict__.items()
            if not key.startswith("_")
        }

    def add_attribute(self, key, value):
        self.__dict__[key] = value

    def save_as_json(self, path):
        """Saves the configuration as a JSON file."""
        with open(path, "w", encoding="utf-8") as config_file:
            json.dump(
                self.to_dict(),
                config_file,
                indent=2,
                sort_keys=True,
            )
            config_file.write("\n")


class ModelSpec(LayerSpec):
    """The top level layer specification."""

    def __init__(self):
        """Initializes the model specification."""
        self._config = self.get_default_config()
        self._files = {}

    @property
    def name(self):
        """The name of the model specification."""
        raise NotImplementedError()

    @property
    def revision(self):
        """The model specification revision.

        This value is incremented each time the weights layout of the model is
        changed (e.g. a weight is renamed).
        """
        return 1

    @property
    def config(self):
        """The model configuration."""
        return self._config

    def get_default_config(self):
        """Returns the default configuration used by this model."""
        return None

    def register_file(self, path: str, filename: Optional[str] = None) -> None:
        """Registers a file to be saved in the model directory."""
        if not os.path.isfile(path):
            raise ValueError("File %s does not exist" % path)
        if filename is None:
            filename = os.path.basename(path)
        if filename in self._files:
            raise ValueError("A file with name %s was already registered" % filename)
        self._files[filename] = path

    def save(self, output_dir: str) -> None:
        """Saves this model on disk.

        Arguments:
          output_dir: Output directory where the model is saved.
        """
        self._serialize(os.path.join(output_dir, "model.bin"))
        if self._config is not None:
            self._config.save_as_json(os.path.join(output_dir, "config.json"))

        for filename, path in self._files.items():
            destination = os.path.join(output_dir, filename)
            if os.path.exists(destination):
                raise RuntimeError(
                    "File %s already exists in the model directory" % destination
                )
            shutil.copy(path, destination)

    def _serialize(self, path):
        """Serializes the model variables."""
        variables = []
        aliases = []
        for variable in self.variables(ordered=True):
            if isinstance(variable[1], str):
                aliases.append(variable)
            else:
                variables.append(variable)

        with open(path, "wb") as model:

            def _write_string(string):
                model.write(struct.pack("H", len(string) + 1))
                model.write(string.encode("utf-8"))
                model.write(struct.pack("B", 0))

            model.write(struct.pack("I", CURRENT_BINARY_VERSION))
            _write_string(self.name)
            model.write(struct.pack("I", self.revision))
            model.write(struct.pack("I", len(variables)))
            for name, value in variables:
                _write_string(name)
                model.write(struct.pack("B", len(value.shape)))
                for dim in value.shape:
                    model.write(struct.pack("I", dim))
                model.write(struct.pack("B", _dtype_to_type_id(value.dtype)))
                model.write(struct.pack("I", value.num_bytes()))
                model.write(value.to_bytes())
            model.write(struct.pack("I", len(aliases)))
            for alias, variable_name in aliases:
                _write_string(alias)
                _write_string(variable_name)


def _flatten_vocabularies(vocabularies):
    for name, vocabulary in vocabularies.items():
        if len(vocabulary) == 1:
            yield name, vocabulary[0]
        else:
            for i, vocab in enumerate(vocabulary):
                yield "%s_%d" % (name, i + 1), vocab


class SequenceToSequenceModelConfig(ModelConfig):
    """Configuration for sequence-to-sequence models."""

    def __init__(
        self,
        unk_token: str = "<unk>",
        bos_token: str = "<s>",
        eos_token: str = "</s>",
        decoder_start_token: Optional[str] = "<s>",
        add_source_bos: bool = False,
        add_source_eos: bool = False,
        **kwargs,
    ):
        """Initializes the configuration for sequence-to-sequence models.

        Args:
          unk_token: The unknown token.
          bos_token: The start of sentence token.
          eos_token: The end of sentence token.
          decoder_start_token: The decoder start token. If ``None``, the token should
            be passed by the user in the target prefix.
          add_source_bos: If ``True``, ``bos_token`` will be automatically added to
            the source input.
          add_source_eos: If ``True``, ``eos_token`` will be automatically added to
            the source input.
          **kwargs: Additional configuration.
        """
        super().__init__(
            unk_token=unk_token,
            bos_token=bos_token,
            eos_token=eos_token,
            decoder_start_token=decoder_start_token,
            add_source_bos=add_source_bos,
            add_source_eos=add_source_eos,
            **kwargs,
        )


class SequenceToSequenceModelSpec(ModelSpec):
    """Base specification for sequence to sequence models."""

    def __init__(self):
        """Initializes a sequence to sequence model specification."""
        super().__init__()
        self._vocabularies = {
            "source": [],
            "target": [],
        }

    def get_default_config(self):
        return SequenceToSequenceModelConfig()

    @abc.abstractmethod
    def get_source_vocabulary_size(self):
        """Returns the source vocabulary size expected by the model."""
        raise NotImplementedError()

    @abc.abstractmethod
    def get_target_vocabulary_size(self):
        """Returns the target vocabulary size expected by the model."""
        raise NotImplementedError()

    def register_source_vocabulary(self, tokens: List[str]) -> None:
        """Registers a source vocabulary of tokens.

        Arguments:
          tokens: List of source tokens.
        """
        self._vocabularies["source"].append(tokens)

    def register_target_vocabulary(self, tokens: List[str]) -> None:
        """Registers a target vocabulary of tokens.

        Arguments:
          tokens: List of target tokens.
        """
        self._vocabularies["target"].append(tokens)

    def register_vocabulary_mapping(self, path: str) -> None:
        """Registers a vocabulary mapping file.

        Arguments:
          path: Path to the vocabulary mapping file.
        """
        self.register_file(path, "vmap.txt")

    def validate(self) -> None:
        super().validate()

        # Check that vocabularies are registered and have the correct size.
        vocabulary_sizes = {
            "source": self.get_source_vocabulary_size(),
            "target": self.get_target_vocabulary_size(),
        }

        for name, sizes in vocabulary_sizes.items():
            if not isinstance(sizes, list):
                sizes = [sizes]
            vocabularies = self._vocabularies[name]
            if len(vocabularies) != len(sizes):
                raise ValueError(
                    "Incorrect number of %s vocabularies: %d registered, but expected %d"
                    % (name, len(vocabularies), len(sizes))
                )
            for i, (vocabulary, expected_size) in enumerate(zip(vocabularies, sizes)):
                if len(vocabulary) != expected_size:
                    raise ValueError(
                        "%s vocabulary %d has size %d but the model expected a vocabulary "
                        "of size %d"
                        % (name.capitalize(), i, len(vocabulary), expected_size)
                    )

    def save(self, output_dir: str) -> None:
        # Save the vocabularies.
        vocabularies = dict(_flatten_vocabularies(self._vocabularies))
        all_vocabularies = list(vocabularies.values())
        if all(vocabulary == all_vocabularies[0] for vocabulary in all_vocabularies):
            vocabularies = {"shared": all_vocabularies[0]}

        for name, tokens in vocabularies.items():
            _save_vocabulary(output_dir, "%s_vocabulary" % name, tokens)

        # Save the rest of the model.
        super().save(output_dir)


class LanguageModelConfig(ModelConfig):
    """Configuration for language models."""

    def __init__(
        self,
        unk_token: str = "<unk>",
        bos_token: str = "<s>",
        eos_token: str = "</s>",
        **kwargs,
    ):
        """Initializes the configuration for language models.

        Args:
          unk_token: The unknown token.
          bos_token: The start of sentence token.
          eos_token: The end of sentence token.
          **kwargs: Additional configuration.
        """
        super().__init__(
            unk_token=unk_token,
            bos_token=bos_token,
            eos_token=eos_token,
            **kwargs,
        )


class LanguageModelSpec(ModelSpec):
    """Base specification for language models."""

    def __init__(self):
        """Initializes a language model specification."""
        super().__init__()
        self._vocabulary = []

    def get_default_config(self):
        return LanguageModelConfig()

    @abc.abstractmethod
    def get_vocabulary_size(self):
        """Returns the vocabulary size expected by the model."""
        raise NotImplementedError()

    def register_vocabulary(self, tokens: List[str]) -> None:
        """Registers the vocabulary of tokens.

        Arguments:
          tokens: List of tokens.
        """
        self._vocabulary = list(tokens)

    def validate(self) -> None:
        super().validate()

        expected_vocabulary_size = self.get_vocabulary_size()
        if len(self._vocabulary) != expected_vocabulary_size:
            raise ValueError(
                "Vocabulary has size %d but the model expected a vocabulary of size %d"
                % (len(self._vocabulary), expected_vocabulary_size)
            )

    def save(self, output_dir: str) -> None:
        # Save the vocabulary.
        _save_vocabulary(output_dir, "vocabulary", self._vocabulary)

        # Save the rest of the model.
        super().save(output_dir)


def _save_vocabulary(output_dir, name, tokens):
    vocabulary_path = os.path.join(output_dir, "%s.json" % name)

    with open(vocabulary_path, "w", encoding="utf-8") as vocabulary_file:
        json.dump(tokens, vocabulary_file, indent=2)


class Variable(abc.ABC):
    """Abstract base class for model variables."""

    @property
    @abc.abstractmethod
    def shape(self) -> List[int]:
        raise NotImplementedError()

    def is_scalar(self) -> bool:
        return len(self.shape) == 0

    @property
    @abc.abstractmethod
    def dtype(self) -> str:
        raise NotImplementedError()

    def to(self, dtype: str) -> "Variable":
        if dtype == self.dtype:
            return self
        return self._to(dtype)

    @abc.abstractmethod
    def numpy(self) -> np.ndarray:
        raise NotImplementedError()

    def equal(self, other) -> bool:
        return type(self) is type(other) and self._equal(other)

    @abc.abstractmethod
    def num_bytes(self) -> int:
        raise NotImplementedError()

    @abc.abstractmethod
    def to_bytes(self) -> bytes:
        raise NotImplementedError()

    @abc.abstractmethod
    def _to(self, dtype: str) -> "Variable":
        raise NotImplementedError()

    @abc.abstractmethod
    def _equal(self, other) -> bool:
        raise NotImplementedError()


class NumpyVariable(Variable):
    """Model variable as a Numpy array."""

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

    @property
    def shape(self) -> List[int]:
        return self.array.shape

    @property
    def dtype(self) -> str:
        return self.array.dtype.name

    def numpy(self) -> np.ndarray:
        return self.array

    def num_bytes(self) -> int:
        return self.array.nbytes

    def to_bytes(self) -> bytes:
        return self.array.tobytes()

    def _to(self, dtype: str) -> Variable:
        if dtype == "bfloat16":
            if not torch_is_available:
                raise RuntimeError(
                    "Converting to bfloat16 requires torch to be installed"
                )
            return PyTorchVariable.from_numpy(self.array).to(dtype)

        dtype = np.dtype(dtype)
        self.array = self.array.astype(dtype)
        return self

    def _equal(self, other) -> bool:
        a = self.array
        b = other.array
        return a is b or (
            a.dtype == b.dtype
            and a.shape == b.shape
            and a.flat[0] == b.flat[0]
            and np.array_equal(a, b)
        )


class PyTorchVariable(Variable):
    """Model variable as a PyTorch tensor."""

    def __init__(self, tensor):
        if isinstance(tensor, torch.nn.Parameter):
            tensor = tensor.data

        self.tensor = tensor.contiguous()

    @classmethod
    def from_numpy(cls, array):
        tensor = torch.from_numpy(array)
        return cls(tensor)

    @property
    def shape(self) -> List[int]:
        return list(self.tensor.shape)

    @property
    def dtype(self) -> str:
        return str(self.tensor.dtype).replace("torch.", "")

    def numpy(self) -> np.ndarray:
        return self.tensor.detach().numpy()

    def num_bytes(self) -> int:
        return self.tensor.numel() * self.tensor.element_size()

    def to_bytes(self) -> bytes:
        max_size = 2**31 - 1
        num_bytes = self.num_bytes()
        output = b""
        offset = 0
        while num_bytes > 0:
            chunk_size = max_size if num_bytes > max_size else num_bytes
            chunk = ctypes.string_at(self.tensor.data_ptr() + offset, chunk_size)
            output += chunk
            offset += chunk_size
            num_bytes -= chunk_size
        return output

    def _to(self, dtype: str) -> Variable:
        dtype = getattr(torch, dtype)
        self.tensor = self.tensor.to(dtype)
        return self

    def _equal(self, other) -> bool:
        a = self.tensor
        b = other.tensor
        return a is b or (a.dtype == b.dtype and torch.equal(a, b))