import tensorflow as tf
from tensorflow.keras import layers, Model

# Toy interaction data
user_ids = tf.constant([10, 42, 7, 99, 5, 5, 10, 12], dtype=tf.int32)    # shape: (N,)
item_ids = tf.constant([1001, 5, 8, 8, 7, 3, 1001, 5], dtype=tf.int32)   # shape: (N,)
labels   = tf.constant([1, 0, 1, 0, 1, 0, 1, 0], dtype=tf.float32)       # click/purchase

num_users, num_items, emb_dim = 100_000, 50_000, 64

# Two-tower architecture
user_in = layers.Input(shape=(), dtype=tf.int32, name="user_id")          # <-- pass scalars per row
item_in = layers.Input(shape=(), dtype=tf.int32, name="item_id")

u = layers.Embedding(num_users, emb_dim, name="user_emb")(user_in)        # (batch, emb_dim)
v = layers.Embedding(num_items, emb_dim, name="item_emb")(item_in)        # (batch, emb_dim)

score = layers.Dot(axes=1)([u, v])                                        # (batch, 1)
out = layers.Activation("sigmoid")(score)

recsys = Model([user_in, item_in], out)
recsys.compile(optimizer="adam", loss="binary_crossentropy", metrics=["AUC"])

# Train — you can pass a dict keyed by Input names or a list [user_ids, item_ids]
# shapes: user_ids: (batch,), item_ids: (batch,), y: (batch,)
recsys.fit({"user_id": user_ids, "item_id": item_ids}, labels,
           batch_size=4, epochs=5, validation_split=0.25)

# Extract embeddings for similarity search
user_embeddings = recsys.get_layer("user_emb").get_weights()[0]           # (num_users, emb_dim)
item_embeddings = recsys.get_layer("item_emb").get_weights()[0]           # (num_items, emb_dim)

# Find similar items via dot product
def find_similar_items(item_id, top_k=5):
    query_vec = item_embeddings[item_id]                                  # (emb_dim,)
    scores = item_embeddings @ query_vec                                  # (num_items,)
    top_indices = scores.argsort()[-top_k-1:-1][::-1]                    # exclude self
    return top_indices, scores[top_indices]

similar_items, similarity_scores = find_similar_items(item_id=5)
print(f"Items similar to 5: {similar_items}")

        __
        __