import torch
from transformers import BertTokenizer, BertModel
from torch.nn import Embedding
import numpy as np

# BERT 모델 및 토크나이저 로드
tokenizer = BertTokenizer.from_pretrained("klue/bert-base")
bert_model = BertModel.from_pretrained("klue/bert-base")

# 상품 데이터 임베딩
def embed_product_data(product_data):
    text = product_data.get("title", "") + " " + product_data.get("description", "")
    inputs = tokenizer(
        text, return_tensors="pt", truncation=True, padding=True, max_length=128
    )
    outputs = bert_model(**inputs)
    text_embedding = outputs.last_hidden_state.mean(dim=1)

    category_embedding_layer = Embedding(num_embeddings=50, embedding_dim=16)
    color_embedding_layer = Embedding(num_embeddings=20, embedding_dim=8)

    category_id = product_data.get("category_id", 0)
    color_id = product_data.get("color_id", 0)

    category_embedding = category_embedding_layer(torch.tensor([category_id]))
    color_embedding = color_embedding_layer(torch.tensor([color_id]))

    combined_embedding = torch.cat((text_embedding, category_embedding, color_embedding), dim=1)
    return combined_embedding.detach().numpy()

# 사용자 데이터 임베딩
def embed_user_data(user_data):
    embedding_layer = Embedding(num_embeddings=100, embedding_dim=128)

    gender_id = 0 if user_data['gender'] == 'M' else 1
    scaled_height = int((user_data['height'] - 50) * 99 // 200)
    scaled_weight = int((user_data['weight'] - 30) * 99 // 170)

    age_embedding = embedding_layer(torch.tensor([user_data['age']])).view(1, -1)
    gender_embedding = embedding_layer(torch.tensor([gender_id])).view(1, -1)
    height_embedding = embedding_layer(torch.tensor([scaled_height])).view(1, -1)
    weight_embedding = embedding_layer(torch.tensor([scaled_weight])).view(1, -1)

    combined_embedding = torch.cat((age_embedding, gender_embedding, height_embedding, weight_embedding), dim=1)
    return combined_embedding.detach().numpy()