from typing import Any, Callable, Dict, List, Optional, Sequence, Union

import tensorflow as tf

from merlin.models.tf.blocks.mlp import MLPBlock
from merlin.models.tf.blocks.retrieval.matrix_factorization import QueryItemIdsEmbeddingsBlock
from merlin.models.tf.blocks.retrieval.two_tower import TwoTowerBlock
from merlin.models.tf.blocks.sampling.base import ItemSampler
from merlin.models.tf.core.base import Block, BlockType
from merlin.models.tf.core.combinators import ParallelBlock
from merlin.models.tf.core.encoder import EmbeddingEncoder, Encoder
from merlin.models.tf.inputs.base import InputBlock
from merlin.models.tf.inputs.embedding import EmbeddingOptions
from merlin.models.tf.models.base import Model, RetrievalModel, RetrievalModelV2
from merlin.models.tf.models.utils import parse_prediction_tasks
from merlin.models.tf.outputs.base import DotProduct, ModelOutput
from merlin.models.tf.outputs.contrastive import ContrastiveOutput
from merlin.models.tf.outputs.sampling.base import ItemSamplersType
from merlin.models.tf.outputs.sampling.popularity import PopularityBasedSamplerV2
from merlin.models.tf.prediction_tasks.base import ParallelPredictionBlock, PredictionTask
from merlin.models.tf.prediction_tasks.next_item import NextItemPredictionTask
from merlin.models.tf.prediction_tasks.retrieval import ItemRetrievalTask
from merlin.models.utils.schema_utils import categorical_cardinalities
from merlin.schema import Schema, Tags


def MatrixFactorizationModel(
    schema: Schema,
    dim: int,
    query_id_tag=Tags.USER_ID,
    item_id_tag=Tags.ITEM_ID,
    embeddings_initializers: Optional[
        Union[Dict[str, Callable[[Any], None]], Callable[[Any], None]]
    ] = None,
    embeddings_l2_reg: float = 0.0,
    post: Optional[BlockType] = None,
    prediction_tasks: Optional[
        Union[PredictionTask, List[PredictionTask], ParallelPredictionBlock]
    ] = None,
    logits_temperature: float = 1.0,
    samplers: Sequence[ItemSampler] = (),
    **kwargs,
) -> RetrievalModel:
    """Builds a matrix factorization model.

    Example Usage::
        mf = MatrixFactorizationModel(schema, dim=128)
        mf.compile(optimizer="adam")
        mf.fit(train_data, epochs=10)

    Parameters
    ----------
    schema: Schema
        The `Schema` with the input features
    dim: int
        The dimension of the embeddings.
    query_id_tag : Tag
        The tag to select query features, by default `Tags.USER`
    item_id_tag : Tag
        The tag to select item features, by default `Tags.ITEM`
    embeddings_initializers : Optional[Dict[str, Callable[[Any], None]]] = None
        An initializer function or a dict where keys are feature names and values are
        callable to initialize embedding tables
    embeddings_l2_reg: float = 0.0
        Factor for L2 regularization of the embeddings vectors (from the current batch only)
    post: Optional[Block], optional
        The optional `Block` to apply on both outputs of Two-tower model
    prediction_tasks: optional
        The optional `PredictionTask` or list of `PredictionTask` to apply on the model.
    logits_temperature: float
        Parameter used to reduce model overconfidence, so that logits / T.
        Defaults to 1.
    loss: Optional[LossType]
        Loss function.
        Defaults to `bpr`.
    samplers: List[ItemSampler]
        List of samplers for negative sampling, by default `[InBatchSampler()]`

    Returns
    -------
    RetrievalModel
    """

    if not prediction_tasks:
        prediction_tasks = ItemRetrievalTask(
            schema,
            logits_temperature=logits_temperature,
            samplers=list(samplers),
            **kwargs,
        )

    prediction_tasks = parse_prediction_tasks(schema, prediction_tasks)
    mf = QueryItemIdsEmbeddingsBlock(
        schema=schema,
        dim=dim,
        query_id_tag=query_id_tag,
        item_id_tag=item_id_tag,
        embeddings_initializers=embeddings_initializers,
        embeddings_l2_reg=embeddings_l2_reg,
        post=post,
        **kwargs,
    )

    model = RetrievalModel(mf, prediction_tasks)

    return model


def TwoTowerModel(
    schema: Schema,
    query_tower: Block,
    item_tower: Optional[Block] = None,
    query_tower_tag=Tags.USER,
    item_tower_tag=Tags.ITEM,
    embedding_options: EmbeddingOptions = EmbeddingOptions(
        embedding_dims=None,
        embedding_dim_default=64,
        infer_embedding_sizes=False,
        infer_embedding_sizes_multiplier=2.0,
    ),
    post: Optional[BlockType] = None,
    prediction_tasks: Optional[
        Union[PredictionTask, List[PredictionTask], ParallelPredictionBlock]
    ] = None,
    logits_temperature: float = 1.0,
    samplers: Sequence[ItemSampler] = (),
    **kwargs,
) -> RetrievalModel:
    """Builds the Two-tower architecture, as proposed in [1].

    Example Usage::
        two_tower = TwoTowerModel(schema, MLPBlock([256, 64]))
        two_tower.compile(optimizer="adam")
        two_tower.fit(train_data, epochs=10)

    References
    ----------
    [1] Yi, Xinyang, et al.
        "Sampling-bias-corrected neural modeling for large corpus item recommendations."
        Proceedings of the 13th ACM Conference on Recommender Systems. 2019.

    Parameters
    ----------
    schema: Schema
        The `Schema` with the input features
    query_tower: Block
        The `Block` that combines user features
    item_tower: Optional[Block], optional
        The optional `Block` that combines items features.
        If not provided, a copy of the query_tower is used.
    query_tower_tag: Tag
        The tag to select query features, by default `Tags.USER`
    item_tower_tag: Tag
        The tag to select item features, by default `Tags.ITEM`
    embedding_options : EmbeddingOptions
        Options for the input embeddings.
        - embedding_dims: Optional[Dict[str, int]] - The dimension of the
        embedding table for each feature (key), by default {}
        - embedding_dim_default: int - Default dimension of the embedding
        table, when the feature is not found in ``embedding_dims``, by default 64
        - infer_embedding_sizes : bool, Automatically defines the embedding
        dimension from the feature cardinality in the schema, by default False
        - infer_embedding_sizes_multiplier: int. Multiplier used by the heuristic
        to infer the embedding dimension from its cardinality. Generally
        reasonable values range between 2.0 and 10.0. By default 2.0.
    post: Optional[Block], optional
        The optional `Block` to apply on both outputs of Two-tower model
    prediction_tasks: optional
        The optional `PredictionTask` or list of `PredictionTask` to apply on the model.
    logits_temperature: float
        Parameter used to reduce model overconfidence, so that logits / T.
        Defaults to 1.
    loss: Optional[LossType]
        Loss function.
        Defaults to `categorical_crossentropy`.
    samplers: List[ItemSampler]
        List of samplers for negative sampling, by default `[InBatchSampler()]`

    Returns
    -------
    RetrievalModel
    """

    if not prediction_tasks:
        prediction_tasks = ItemRetrievalTask(
            schema,
            logits_temperature=logits_temperature,
            samplers=list(samplers),
            **kwargs,
        )

    prediction_tasks = parse_prediction_tasks(schema, prediction_tasks)
    two_tower = TwoTowerBlock(
        schema=schema,
        query_tower=query_tower,
        item_tower=item_tower,
        query_tower_tag=query_tower_tag,
        item_tower_tag=item_tower_tag,
        embedding_options=embedding_options,
        post=post,
        **kwargs,
    )

    model = RetrievalModel(two_tower, prediction_tasks)

    return model


def YoutubeDNNRetrievalModel(
    schema: Schema,
    aggregation: str = "concat",
    top_block: Block = MLPBlock([64]),
    l2_normalization: bool = True,
    extra_pre_call: Optional[Block] = None,
    task_block: Optional[Block] = None,
    logits_temperature: float = 1.0,
    sampled_softmax: bool = True,
    num_sampled: int = 100,
    min_sampled_id: int = 0,
    embedding_options: EmbeddingOptions = EmbeddingOptions(
        embedding_dims=None,
        embedding_dim_default=64,
        infer_embedding_sizes=False,
        infer_embedding_sizes_multiplier=2.0,
    ),
) -> Model:
    """Build the Youtube-DNN retrieval model.
    More details of the architecture can be found in [1]_.
    The sampled_softmax is enabled by default [2]_ [3]_ [4]_.

    Example Usage::
        model = YoutubeDNNRetrievalModel(schema, num_sampled=100)
        model.compile(optimizer="adam")
        model.fit(train_data, epochs=10)

    References
    ----------
    .. [1] Covington, Paul, Jay Adams, and Emre Sargin.
        "Deep neural networks for youtube recommendations."
        Proceedings of the 10th ACM conference on recommender systems. 2016.

    .. [2] Yoshua Bengio and Jean-Sébastien Sénécal. 2003. Quick Training of Probabilistic
       Neural Nets by Importance Sampling. In Proceedings of the conference on Artificial
       Intelligence and Statistics (AISTATS).

    .. [3] Y. Bengio and J. S. Senecal. 2008. Adaptive Importance Sampling to Accelerate
       Training of a Neural Probabilistic Language Model. Trans. Neur. Netw. 19, 4 (April
       2008), 713–722. https://doi.org/10.1109/TNN.2007.912312

    .. [4] Jean, Sébastien, et al. "On using very large target vocabulary for neural
        machine translation." arXiv preprint arXiv:1412.2007 (2014).


    Parameters
    ----------
    schema: Schema
        The `Schema` with the input features
    aggregation: str
        The aggregation method to use for the sequence of features.
        Defaults to `concat`.
    top_block: Block
        The `Block` that combines the top features
    l2_normalization: bool
        Whether to apply L2 normalization before computing dot interactions.
        Defaults to True.
    extra_pre_call: Optional[Block]
        The optional `Block` to apply before the model.
    task_block: Optional[Block]
        The optional `Block` to apply on the model.
    logits_temperature: float
        Parameter used to reduce model overconfidence, so that logits / T.
        Defaults to 1.
    sampled_softmax: bool
        Compute the logits scores over all items of the catalog or
        generate a subset of candidates
        Defaults to False
    num_sampled: int
        When sampled_softmax is enabled, specify the number of
        negative candidates to generate for each batch.
        Defaults to 100
    min_sampled_id: int
        The minimum id value to be sampled with sampled softmax.
        Useful to ignore the first categorical
        encoded ids, which are usually reserved for <nulls>,
        out-of-vocabulary or padding. Defaults to 0.
    embedding_options : EmbeddingOptions, optional
        An EmbeddingOptions instance, which allows for a number of
        options for the embedding table, by default EmbeddingOptions()
    """

    inputs = InputBlock(
        schema,
        aggregation=aggregation,
        embedding_options=embedding_options,
    )

    task = NextItemPredictionTask(
        schema=schema,
        weight_tying=False,
        extra_pre_call=extra_pre_call,
        task_block=task_block,
        logits_temperature=logits_temperature,
        l2_normalization=l2_normalization,
        sampled_softmax=sampled_softmax,
        num_sampled=num_sampled,
        min_sampled_id=min_sampled_id,
    )

    # TODO: Figure out how to make this fit as
    # a RetrievalModel (which must have a RetrievalBlock)
    return Model(inputs, top_block, task)


def MatrixFactorizationModelV2(
    schema: Schema,
    dim: int,
    query_id_tag=Tags.USER_ID,
    candidate_id_tag=Tags.ITEM_ID,
    embeddings_initializers: Optional[
        Union[Dict[str, Callable[[Any], None]], Callable[[Any], None]]
    ] = None,
    embeddings_l2_batch_regularization: Union[float, Dict[str, float]] = 0.0,
    post: Optional[BlockType] = None,
    outputs: Optional[Union[ModelOutput, List[ModelOutput]]] = None,
    negative_samplers: ItemSamplersType = None,
    logits_temperature: float = 1.0,
    **kwargs,
) -> RetrievalModelV2:
    """Builds a matrix factorization (MF) model.

    Example Usage::
        mf = MatrixFactorizationModelV2(schema, dim=128)
        mf.compile(optimizer="adam")
        mf.fit(train_data, epochs=10)

    Parameters
    ----------
    schema: Schema
        The `Schema` with the input features
    dim: int
        The dimension of the embeddings.
    query_id_tag : Tag
        The tag to select query-id feature, by default `Tags.USER_ID`
    candidate_id_tag : Tag
        The tag to select cnadidate-id feature, by default `Tags.ITEM_ID`
    embeddings_initializers : Optional[Dict[str, Callable[[Any], None]]] = None
        An initializer function or a dict where keys are feature names and values are
        callable to initialize embedding tables
    embeddings_l2_batch_regularization: Union[float, Dict[str, float]] = 0.0
        Factor for L2 regularization of the embeddings vectors (from the current batch only)
        If a dictionary is provided, the keys are the query-id and candidate-id names and
        the values are the regularization factor.
    post: Optional[Block], optional
        The optional `Block` to apply on both outputs of the MF towers.
    outputs: optional
        The optional `ModelOutput` or list of `ModelOutput` to apply on the MF model.
    negative_samplers: List[ItemSampler]
        List of samplers for negative sampling, by default None
        If the `outputs` and `negative_samplers` are not specified the Matrix Factorization model
        is trained with contrastive learning and `in-batch` negative sampling strategy.
    logits_temperature: float
        Parameter used to reduce model overconfidence, so that logits / T.
        Defaults to 1.


    Returns
    -------
    RetrievalModelV2
    """

    query = schema.select_by_tag(query_id_tag)
    candidate = schema.select_by_tag(candidate_id_tag)

    query_encoder = EmbeddingEncoder(
        query,
        dim=dim,
        embeddings_initializer=embeddings_initializers,
        embeddings_l2_batch_regularization=embeddings_l2_batch_regularization,
        post=post,
    )
    candidate_encoder = EmbeddingEncoder(
        candidate,
        dim=dim,
        embeddings_initializer=embeddings_initializers,
        embeddings_l2_batch_regularization=embeddings_l2_batch_regularization,
        post=post,
    )

    if not outputs:
        if not negative_samplers:
            negative_samplers = ["in-batch"]
        outputs = ContrastiveOutput(
            to_call=DotProduct(),
            negative_samplers=negative_samplers,
            logits_temperature=logits_temperature,
            schema=candidate,
            **kwargs,
        )

    if isinstance(outputs, list):
        outputs = ParallelBlock(*outputs)

    model = RetrievalModelV2(
        query=query_encoder,
        candidate=candidate_encoder,
        output=outputs,
    )

    return model


def TwoTowerModelV2(
    query_tower: Encoder,
    candidate_tower: Encoder,
    candidate_id_tag=Tags.ITEM_ID,
    outputs: Optional[Union[ModelOutput, List[ModelOutput]]] = None,
    logits_temperature: float = 1.0,
    negative_samplers: ItemSamplersType = None,
    **kwargs,
) -> RetrievalModelV2:
    """Builds the Two-tower architecture, as proposed in [1].

    Example Usage::
        query = mm.Encoder(user_schema, mm.MLPBlock([128]))
        candidate = mm.Encoder(item_schema, mm.MLPBlock([128]))
        model = TwoTowerModel(query, candidate)
        two_tower.compile(optimizer="adam")
        two_tower.fit(train_data, epochs=10)

    References
    ----------
    [1] Yi, Xinyang, et al.
        "Sampling-bias-corrected neural modeling for large corpus item recommendations."
        Proceedings of the 13th ACM Conference on Recommender Systems. 2019.

    Parameters
    ----------
    query_tower: Encoder
        The layer that encodes query features
    candidate_tower: Encoder
        The  layer that encodes candidates features
    candidate_id_tag: Tag, optional
        The tag to select candidate-id feature, by default `Tags.ITEM_ID`
    outputs:  Union[ModelOutput, List[ModelOutput]], optional
        The optional `ModelOutput` or list of `ModelOutput` to apply on the model.
    logits_temperature: float
        Parameter used to reduce model overconfidence, so that logits / T.
        Defaults to 1.
    negative_samplers: List[ItemSampler]
        List of samplers for negative sampling, by default None
        If the `outputs` and `negative_samplers` are not specified the two tower model
        is trained with contrastive learning and `in-batch` negative sampling strategy.

    Returns
    -------
    RetrievalModelV2
    """
    assert isinstance(query_tower, Encoder), ValueError(
        "The query tower should be an instance of `Encoder` class"
    )
    assert isinstance(candidate_tower, Encoder), ValueError(
        "The query tower should be an instance of `Encoder` class"
    )

    if not outputs:
        if not negative_samplers:
            negative_samplers = ["in-batch"]
        outputs = ContrastiveOutput(
            to_call=DotProduct(),
            negative_samplers=negative_samplers,
            logits_temperature=logits_temperature,
            schema=candidate_tower._schema.select_by_tag(candidate_id_tag),
            **kwargs,
        )

    if isinstance(outputs, list):
        outputs = ParallelBlock(*outputs)

    model = RetrievalModelV2(
        query=query_tower,
        candidate=candidate_tower,
        output=outputs,
    )

    return model


def YoutubeDNNRetrievalModelV2(
    schema: Schema,
    candidate_id_tag=Tags.ITEM_ID,
    top_block: Optional[tf.keras.layers.Layer] = MLPBlock([64]),
    post: Optional[tf.keras.layers.Layer] = None,
    inputs: tf.keras.layers.Layer = None,
    outputs: Optional[Union[ModelOutput, List[ModelOutput]]] = None,
    logits_temperature: float = 1.0,
    num_sampled: int = 100,
    min_sampled_id: int = 0,
    **kwargs,
) -> RetrievalModelV2:
    """Build the Youtube-DNN retrieval model.
    More details of the architecture can be found in [1]_.
    Training with sampled_softmax is enabled by default [2]_ [3]_ [4]_.

    Example Usage::
        model = YoutubeDNNRetrievalModelV2(schema, num_sampled=100)
        model.compile(optimizer="adam")
        model.fit(train_data, epochs=10)

    References
    ----------
    .. [1] Covington, Paul, Jay Adams, and Emre Sargin.
        "Deep neural networks for youtube recommendations."
        Proceedings of the 10th ACM conference on recommender systems. 2016.

    .. [2] Yoshua Bengio and Jean-Sébastien Sénécal. 2003. Quick Training of Probabilistic
       Neural Nets by Importance Sampling. In Proceedings of the conference on Artificial
       Intelligence and Statistics (AISTATS).

    .. [3] Y. Bengio and J. S. Senecal. 2008. Adaptive Importance Sampling to Accelerate
       Training of a Neural Probabilistic Language Model. Trans. Neur. Netw. 19, 4 (April
       2008), 713–722. https://doi.org/10.1109/TNN.2007.912312

    .. [4] Jean, Sébastien, et al. "On using very large target vocabulary for neural
        machine translation." arXiv preprint arXiv:1412.2007 (2014).


    Parameters
    ----------
    schema: Schema
        The `Schema` with the input features
    candidate_id_tag : Tag
        The tag to select cnadidate-id feature, by default `Tags.ITEM_ID`
    top_block: tf.keras.layers.Layer
        The hidden layers to apply on top of the features representation
        vector.
    inputs: tf.keras.layers.Layer, optional
        The input layer to encode input features (sparse and context features)
        If not specified, the input layer is inferred from the schema
        By default None
    post: Optional[tf.keras.layers.Layer], optional
        The optional layer to apply on top of the query encoder.
    logits_temperature: float, optional
        Parameter used to reduce model overconfidence, so that logits / T.
        Defaults to 1.
    num_sampled: int, optional
        When sampled_softmax is enabled, specify the number of
        negative candidates to generate for each batch.
        By default 100
    min_sampled_id: int, optional
        The minimum id value to be sampled with sampled softmax.
        Useful to ignore the first categorical
        encoded ids, which are usually reserved for <nulls>,
        out-of-vocabulary or padding.
        By default 0.
    """
    if not inputs:
        inputs = schema

    candidate = schema.select_by_tag(candidate_id_tag)
    if not candidate:
        raise ValueError(f"The schema should contain a feature tagged as `{candidate_id_tag}`")

    query = Encoder(inputs, top_block, post=post)

    if not outputs:
        cardinalities = categorical_cardinalities(candidate)
        candidate = candidate.first
        candidate_table = query.first["categorical"][candidate.name]
        num_classes = cardinalities[candidate.name]

        outputs = ContrastiveOutput(
            to_call=candidate_table,
            logits_temperature=logits_temperature,
            post=post,
            negative_samplers=PopularityBasedSamplerV2(
                max_num_samples=num_sampled, max_id=num_classes - 1, min_id=min_sampled_id
            ),
        )

    return RetrievalModelV2(query=query, output=outputs)
