from typing import List, Optional
import tml.core.config as base_config
from tml.optimizers import config as optimizer_config

import pydantic


class EmbeddingSnapshot(base_config.BaseConfig):
  """Configuration for Embedding snapshot"""

  emb_name: str = pydantic.Field(
    ..., description="Name of the embedding table from the loaded snapshot"
  )
  embedding_snapshot_uri: str = pydantic.Field(
    ..., description="Path to torchsnapshot of the embedding"
  )


# https://pytorch.org/torchrec/torchrec.modules.html#torchrec.modules.embedding_configs.EmbeddingBagConfig
class EmbeddingBagConfig(base_config.BaseConfig):
  """Configuration for EmbeddingBag."""

  name: str = pydantic.Field(..., description="name of embedding bag")
  num_embeddings: int = pydantic.Field(..., description="size of embedding dictionary")
  embedding_dim: int = pydantic.Field(..., description="size of each embedding vector")
  pretrained: EmbeddingSnapshot = pydantic.Field(None, description="Snapshot properties")
  vocab: str = pydantic.Field(
    None, description="Directory to parquet files of mapping from entity ID to table index."
  )


class EmbeddingOptimizerConfig(base_config.BaseConfig):
  learning_rate: optimizer_config.LearningRate = pydantic.Field(
    None, description="learning rate scheduler for the EBC"
  )
  init_learning_rate: float = pydantic.Field(description="initial learning rate for the EBC")
  # NB: Only sgd is supported right now and implicitly.
  # FBGemm only supports simple exact_sgd which only takes LR as an argument.


class LargeEmbeddingsConfig(base_config.BaseConfig):
  """Configuration for EmbeddingBagCollection.

  The tables listed in this config are gathered into a single torchrec EmbeddingBagCollection.
  """

  tables: List[EmbeddingBagConfig] = pydantic.Field(..., description="list of embedding tables")
  optimizer: EmbeddingOptimizerConfig
  tables_to_log: List[str] = pydantic.Field(
    None, description="list of embedding table names that we want to log during training"
  )


class StratifierConfig(base_config.BaseConfig):
  name: str
  index: int
  value: int


class SmallEmbeddingBagConfig(base_config.BaseConfig):
  """Configuration for SmallEmbeddingBag."""

  name: str = pydantic.Field(..., description="name of embedding bag")
  num_embeddings: int = pydantic.Field(..., description="size of embedding dictionary")
  embedding_dim: int = pydantic.Field(..., description="size of each embedding vector")
  index: int = pydantic.Field(..., description="index in the discrete tensor to look for")


class SmallEmbeddingBagConfig(base_config.BaseConfig):
  """Configuration for SmallEmbeddingBag."""

  name: str = pydantic.Field(..., description="name of embedding bag")
  num_embeddings: int = pydantic.Field(..., description="size of embedding dictionary")
  embedding_dim: int = pydantic.Field(..., description="size of each embedding vector")
  index: int = pydantic.Field(..., description="index in the discrete tensor to look for")


class SmallEmbeddingsConfig(base_config.BaseConfig):
  """Configuration for SmallEmbeddingConfig.

  Here we can use discrete features that already are present in our TFRecords generated using
  segdense conversion as "home_recap_2022_discrete__segdense_vals" which are available in
  the model as "discrete_features", and embed a user-defined set of them with configurable
  dimensions and vocabulary sizes.

  Compared with LargeEmbedding, this config is for small embedding tables that can fit inside
  the model, whereas LargeEmbedding usually is meant to be hydrated outside the model at
  serving time due to size (>>1 GB).

  This small embeddings table uses the same optimizer as the rest of the model."""

  tables: List[SmallEmbeddingBagConfig] = pydantic.Field(
    ..., description="list of embedding tables"
  )