import math
import torch
import logging

import numpy as np
import torch.nn as nn

from typing import *
from abc import ABCMeta
from collections import defaultdict
from cfdata.tabular.misc import np_int_type

from ..protocol import DataLoaderProtocol
from ..misc.toolkit import to_torch
from ..misc.toolkit import Lambda
from ..misc.toolkit import Initializer
from ..misc.toolkit import LoggingMixinWithRank
from ..modules.auxiliary import Dropout


class EncodingResult(NamedTuple):
    one_hot: Optional[torch.Tensor]
    embedding: Optional[torch.Tensor]

    @property
    def merged(self) -> torch.Tensor:
        if self.one_hot is None and self.embedding is None:
            raise ValueError("no data is provided in `EncodingResult`")
        if self.one_hot is None:
            assert self.embedding is not None
            return self.embedding
        if self.embedding is None:
            assert self.one_hot is not None
            return self.one_hot
        return torch.cat([self.one_hot, self.embedding], dim=1)


class OneHot(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        one_hot_fn = lambda column: nn.functional.one_hot(column, dim)
        self.core = Lambda(one_hot_fn, f"one_hot_{dim}")
        self.dim = dim

    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
        return self.core(tensor).to(torch.float32)


class Embedding(nn.Module):
    def __init__(
        self,
        in_dim: int,
        out_dim: int,
        init_method: Optional[str],
        init_config: Dict[str, Any],
    ):
        super().__init__()
        weights = torch.empty(in_dim, out_dim)
        if init_method is None:
            nn.init.normal_(weights)
        else:
            Initializer(init_config).initialize(weights, init_method)
        self.weights = nn.Parameter(weights)
        embedding_fn = lambda column: nn.functional.embedding(column, self.weights)
        self.core = Lambda(embedding_fn, f"embedding: {in_dim} -> {out_dim}")
        self.in_dim, self.out_dim = in_dim, out_dim

    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
        return self.core(tensor)


class Encoder(nn.Module, LoggingMixinWithRank, metaclass=ABCMeta):
    def __init__(
        self,
        config: Dict[str, Any],
        input_dims: List[int],
        methods_list: List[Union[str, List[str]]],
        configs: List[Dict[str, Any]],
        categorical_columns: List[int],
        loaders: Dict[str, DataLoaderProtocol],
    ):
        super().__init__()
        self._fe_init_method: Optional[str]
        self._fe_init_config: Optional[Dict[str, Any]]
        self._init_config(config)
        for loader in loaders.values():
            if loader.enabled_sampling:
                raise ValueError("`loader` should not enable sampling in `Encoder`")
            if loader.sampler.shuffle:
                raise ValueError("`loader` should not be shuffled in `Encoder`")
        self.merged_dim = 0
        self.one_hot_dim = 0
        self.embedding_dim = 0
        dims_tensor = torch.tensor(input_dims, dtype=torch.float32)
        self.register_buffer("input_dims", dims_tensor)
        self.tgt_columns = np.array(sorted(categorical_columns), np_int_type)
        self.merged_dims: Dict[int, int] = defaultdict(int)
        self.embeddings = nn.ModuleList()
        self.one_hot_encoders = nn.ModuleList()
        self._one_hot_indices: List[int] = []
        self._embed_indices: List[int] = []
        self._embed_dims: List[int] = []
        for i, (in_dim, methods, config) in enumerate(
            zip(input_dims, methods_list, configs)
        ):
            if isinstance(methods, str):
                methods = [methods]
            self._register(i, in_dim, methods, config)
        # fast embedding
        if self.use_embedding and self._use_fast_embed:
            if isinstance(self._unified_embed_dim, int):
                unified_embed_dim = self._unified_embed_dim
            else:
                if self._unified_embed_dim == "max":
                    unified_embed_dim = max(self._embed_dims)
                elif self._unified_embed_dim == "mean":
                    unified_embed_dim = int(
                        round(sum(self._embed_dims) / self.num_embedding)
                    )
                elif self._unified_embed_dim == "median":
                    half_idx = self.num_embedding // 2
                    sorted_dims = sorted(self._embed_dims)
                    if self.num_embedding % 2 != 0:
                        unified_embed_dim = sorted_dims[half_idx]
                    else:
                        left = sorted_dims[half_idx - 1]
                        right = sorted_dims[half_idx]
                        unified_embed_dim = int(round(0.5 * (left + right)))
                else:
                    raise NotImplementedError(
                        f"unified embedding dim '{self._unified_embed_dim}' "
                        "is not recognized"
                    )
            for i in self._embed_indices:
                self.merged_dims[i] += unified_embed_dim
            embedding_dim_sum = unified_embed_dim * self.num_embedding
            self.embedding_dim = embedding_dim_sum
            self.merged_dim += embedding_dim_sum
            if self._fe_init_method is None:
                self._fe_init_method = self._de_init_method
            if self._fe_init_config is None:
                self._fe_init_config = self._de_init_config
            assert isinstance(self._fe_init_config, dict)
            self.embeddings.append(
                Embedding(
                    sum(input_dims),
                    unified_embed_dim,
                    self._fe_init_method,
                    self._fe_init_config,
                )
            )
            assert isinstance(self.input_dims, torch.Tensor)
            embed_dims_cumsum = self.input_dims[self._embed_indices].cumsum(0)[:-1]
            self.register_buffer("embed_dims_cumsum", embed_dims_cumsum)
        # embedding dropout
        self.embedding_dropout = None
        if self.use_embedding and 0.0 < self._embed_drop < 1.0:
            self.embedding_dropout = Dropout(self._embed_drop)
        # compile
        self.one_hot_columns = self.tgt_columns[self._one_hot_indices]
        self.embedding_columns = self.tgt_columns[self._embed_indices]
        self._all_one_hot = len(self.one_hot_columns) == len(input_dims)
        self._all_embedding = len(self.embedding_columns) == len(input_dims)
        self._compile(loaders)

    @property
    def num_one_hot(self) -> int:
        return len(self._one_hot_indices)

    @property
    def num_embedding(self) -> int:
        return len(self._embed_indices)

    @property
    def use_one_hot(self) -> bool:
        return self.num_one_hot > 0

    @property
    def use_embedding(self) -> bool:
        return self.num_embedding > 0

    def forward(
        self,
        x_batch: torch.Tensor,
        batch_indices: Optional[np.ndarray],
        loader_name: Optional[str],
    ) -> EncodingResult:
        keys = None
        if loader_name is not None:
            keys = self._get_cache_keys(loader_name)
        categorical_columns = x_batch[..., self.tgt_columns]
        if batch_indices is None or loader_name is None:
            self._oob_imputation(categorical_columns)
        use_cache = keys is not None and batch_indices is not None
        # one hot
        if not self.use_one_hot:
            one_hot = None
        else:
            if use_cache:
                one_hot = getattr(self, keys["one_hot"])[batch_indices]  # type: ignore
            else:
                one_hot_columns = categorical_columns
                if not self._all_one_hot:
                    one_hot_columns = one_hot_columns[..., self._one_hot_indices]
                one_hot = self._one_hot(one_hot_columns)
        # embedding
        if not self.use_embedding:
            embedding = None
        else:
            if not use_cache:
                indices = categorical_columns
            else:
                indices = getattr(self, keys["indices"])[batch_indices]  # type: ignore
            if not self._all_embedding:
                indices = indices[..., self._embed_indices]
            if not use_cache and self._use_fast_embed:
                indices[..., 1:] += self.embed_dims_cumsum
                indices = indices.to(torch.long)
            embedding = self._embedding(indices)
            if self.embedding_dropout is not None:
                embedding = self.embedding_dropout(embedding)
        return EncodingResult(one_hot, embedding)

    def _init_config(self, config: Dict[str, Any]) -> None:
        self.config = config
        self._embed_drop = config.setdefault("embedding_dropout", 0.2)
        self._de_init_method = config.setdefault(
            "default_embedding_init_method", "truncated_normal"
        )
        self._de_init_config = config.setdefault(
            "default_embedding_init_config", {"mean": 0.0, "std": 0.02}
        )
        self._use_fast_embed = config.setdefault("use_fast_embedding", True)
        # [ mean | median | max | int ]
        self._unified_embed_dim = config.setdefault("unified_embedding_dim", "max")
        self._fe_init_method = config.setdefault("fast_embedding_init_method", None)
        self._fe_init_config = config.setdefault("fast_embedding_init_config", None)

    def _register(
        self,
        i: int,
        in_dim: int,
        methods: List[str],
        config: Dict[str, Any],
    ) -> None:
        for method in methods:
            attr = getattr(self, f"_register_{method}", None)
            if attr is None:
                msg = f"encoding method '{method}' is not implemented"
                raise NotImplementedError(msg)
            attr(i, in_dim, config)

    def _register_one_hot(self, i: int, in_dim: int, _: Dict[str, Any]) -> None:
        self.one_hot_encoders.append(OneHot(in_dim))
        self._one_hot_indices.append(i)
        self.merged_dims[i] += in_dim
        self.one_hot_dim += in_dim
        self.merged_dim += in_dim

    @staticmethod
    def _get_embed_key(num: int) -> str:
        return f"embedding_weight_{num}"

    def _register_embedding(self, i: int, in_dim: int, config: Dict[str, Any]) -> None:
        self._embed_indices.append(i)
        embedding_dim = config.setdefault("embedding_dim", "auto")
        if isinstance(embedding_dim, int):
            out_dim = embedding_dim
        elif embedding_dim == "log":
            out_dim = math.ceil(math.log2(in_dim))
        elif embedding_dim == "sqrt":
            out_dim = math.ceil(math.sqrt(in_dim))
        elif embedding_dim == "auto":
            out_dim = min(in_dim, max(4, min(8, math.ceil(math.log2(in_dim)))))
        else:
            raise ValueError(f"embedding dim '{embedding_dim}' is not defined")
        if self._use_fast_embed:
            self._embed_dims.append(out_dim)
            if self._fe_init_method is None:
                self._fe_init_method = config.get("init_method")
            if self._fe_init_config is None:
                self._fe_init_config = config.get("init_config")
        else:
            init_method = config.setdefault("init_method", self._de_init_method)
            init_config = config.setdefault("init_config", self._de_init_config)
            self.merged_dims[i] += out_dim
            self.embedding_dim += out_dim
            self.merged_dim += out_dim
            self.embeddings.append(Embedding(in_dim, out_dim, init_method, init_config))

    def _oob_imputation(
        self,
        categorical_columns: torch.Tensor,
        batch_indices: Optional[np.ndarray] = None,
        loader_name: Optional[str] = None,
    ) -> None:
        if loader_name is None or batch_indices is None:
            oob_mask = categorical_columns >= self.input_dims
        else:
            keys = self._get_cache_keys(loader_name)
            oob_mask = getattr(self, keys["oob"])[batch_indices]
        if torch.any(oob_mask):
            self.log_msg(  # type: ignore
                "out of bound occurred, "
                f"ratio : {torch.mean(oob_mask.to(torch.float32)).item():8.6f}",
                prefix=self.warning_prefix,  # type: ignore
                verbose_level=5,
                msg_level=logging.WARNING,
            )
            # TODO : currently pytorch does not support onnx with bool masks
            #        in the future this line should be un-indented
            categorical_columns[oob_mask] = 0.0

    @staticmethod
    def _to_split(columns: torch.Tensor) -> List[torch.Tensor]:
        return list(columns.to(torch.long).t().unbind())

    def _one_hot(self, one_hot_columns: torch.Tensor) -> torch.Tensor:
        split = self._to_split(one_hot_columns)
        encodings = [
            encoder(flat_feature)
            for encoder, flat_feature in zip(self.one_hot_encoders, split)
        ]
        return torch.cat(encodings, dim=1)

    def _embedding(self, indices_columns: torch.Tensor) -> torch.Tensor:
        if self._use_fast_embed:
            embed_mat = self.embeddings[0](indices_columns)
            return embed_mat.view(-1, self.embedding_dim)
        split = self._to_split(indices_columns)
        encodings = [
            embedding(flat_feature)
            for embedding, flat_feature in zip(self.embeddings, split)
        ]
        return torch.cat(encodings, dim=1)

    @staticmethod
    def _get_cache_keys(name: str) -> Dict[str, str]:
        return {
            "one_hot": f"{name}_one_hot_cache",
            "indices": f"{name}_indices_cache",
            "oob": f"{name}_oob_cache",
        }

    def _compile(self, loaders: Dict[str, DataLoaderProtocol]) -> None:
        for name, loader in loaders.items():
            categorical_features = []
            return_indices = loader.return_indices
            for sample in loader:
                if return_indices:
                    assert isinstance(sample, tuple)
                    sample = sample[0]
                assert isinstance(sample, dict)
                x_batch = sample["x_batch"]
                categorical_features.append(x_batch[..., self.tgt_columns])
            tensor = to_torch(np.vstack(categorical_features))
            keys = self._get_cache_keys(name)
            # compile oob
            self.register_buffer(keys["oob"], tensor >= self.input_dims)
            self._oob_imputation(
                tensor,
                batch_indices=np.arange(len(tensor)),
                loader_name=name,
            )
            # compile one hot
            if self.use_one_hot:
                one_hot_cache = self._one_hot(tensor[..., self._one_hot_indices])
                self.register_buffer(keys["one_hot"], one_hot_cache)
            # compile embedding
            if self.use_embedding and self._use_fast_embed:
                tensor[..., 1:] += self.embed_dims_cumsum
                indices = tensor.to(torch.long)
                self.register_buffer(keys["indices"], indices)


__all__ = ["Encoder", "EncodingResult"]
