from collections import namedtuple
from functools import partial

import torch
from torchrec.models import deepfm, dlrm
from torchrec.modules.embedding_configs import EmbeddingBagConfig
from torchrec.modules.embedding_modules import EmbeddingBagCollection
from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor

from ..registry import ModelAttribute, model_zoo

BATCH = 2
SHAPE = 10


def gen_kt():
    KT = KeyedTensor(keys=["f1", "f2"], length_per_key=[SHAPE, SHAPE], values=torch.rand((BATCH, 2 * SHAPE)))
    return KT


# KeyedJaggedTensor
def gen_kjt():
    KJT = KeyedJaggedTensor.from_offsets_sync(keys=["f1", "f2"],
                                              values=torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]),
                                              offsets=torch.tensor([0, 2, 4, 6, 8]))
    return KJT


data_gen_fn = lambda: dict(features=torch.rand((BATCH, SHAPE)))


def interaction_arch_data_gen_fn():
    KT = gen_kt()
    return dict(dense_features=torch.rand((BATCH, SHAPE)), sparse_features=KT)


def simple_dfm_data_gen_fn():
    KJT = gen_kjt()
    return dict(dense_features=torch.rand((BATCH, SHAPE)), sparse_features=KJT)


def sparse_arch_data_gen_fn():
    KJT = gen_kjt()
    return dict(features=KJT)


def output_transform_fn(x):
    if isinstance(x, KeyedTensor):
        output = dict()
        for key in x.keys():
            output[key] = x[key]
        return output
    else:
        return dict(output=x)


def output_transform_fn(x):
    if isinstance(x, KeyedTensor):
        output = dict()
        for key in x.keys():
            output[key] = x[key]
        return output
    else:
        return dict(output=x)


def get_ebc():
    # EmbeddingBagCollection
    eb1_config = EmbeddingBagConfig(name="t1", embedding_dim=SHAPE, num_embeddings=SHAPE, feature_names=["f1"])
    eb2_config = EmbeddingBagConfig(name="t2", embedding_dim=SHAPE, num_embeddings=SHAPE, feature_names=["f2"])
    return EmbeddingBagCollection(tables=[eb1_config, eb2_config], device=torch.device('cpu'))


def sparse_arch_model_fn():
    ebc = get_ebc()
    return deepfm.SparseArch(ebc)


def simple_deep_fmnn_model_fn():
    ebc = get_ebc()
    return deepfm.SimpleDeepFMNN(SHAPE, ebc, SHAPE, SHAPE)


def dlrm_model_fn():
    ebc = get_ebc()
    return dlrm.DLRM(ebc, SHAPE, [SHAPE, SHAPE], [5, 1])


def dlrm_sparsearch_model_fn():
    ebc = get_ebc()
    return dlrm.SparseArch(ebc)


model_zoo.register(name='deepfm_densearch',
                   model_fn=partial(deepfm.DenseArch, SHAPE, SHAPE, SHAPE),
                   data_gen_fn=data_gen_fn,
                   output_transform_fn=output_transform_fn)

model_zoo.register(name='deepfm_interactionarch',
                   model_fn=partial(deepfm.FMInteractionArch, SHAPE * 3, ["f1", "f2"], SHAPE),
                   data_gen_fn=interaction_arch_data_gen_fn,
                   output_transform_fn=output_transform_fn)

model_zoo.register(name='deepfm_overarch',
                   model_fn=partial(deepfm.OverArch, SHAPE),
                   data_gen_fn=data_gen_fn,
                   output_transform_fn=output_transform_fn)

model_zoo.register(name='deepfm_simpledeepfmnn',
                   model_fn=simple_deep_fmnn_model_fn,
                   data_gen_fn=simple_dfm_data_gen_fn,
                   output_transform_fn=output_transform_fn)

model_zoo.register(name='deepfm_sparsearch',
                   model_fn=sparse_arch_model_fn,
                   data_gen_fn=sparse_arch_data_gen_fn,
                   output_transform_fn=output_transform_fn)

model_zoo.register(name='dlrm',
                   model_fn=dlrm_model_fn,
                   data_gen_fn=simple_dfm_data_gen_fn,
                   output_transform_fn=output_transform_fn)

model_zoo.register(name='dlrm_densearch',
                   model_fn=partial(dlrm.DenseArch, SHAPE, [SHAPE, SHAPE]),
                   data_gen_fn=data_gen_fn,
                   output_transform_fn=output_transform_fn)

model_zoo.register(name='dlrm_interactionarch',
                   model_fn=partial(dlrm.InteractionArch, 2),
                   data_gen_fn=interaction_arch_data_gen_fn,
                   output_transform_fn=output_transform_fn)

model_zoo.register(name='dlrm_overarch',
                   model_fn=partial(dlrm.OverArch, SHAPE, [5, 1]),
                   data_gen_fn=data_gen_fn,
                   output_transform_fn=output_transform_fn)

model_zoo.register(name='dlrm_sparsearch',
                   model_fn=dlrm_sparsearch_model_fn,
                   data_gen_fn=sparse_arch_data_gen_fn,
                   output_transform_fn=output_transform_fn)