#1. Импорт необходимых библиотек
#Сначала нам нужно импортировать необходимые библиотеки.

import open_clip
import torch
from transformers import MarianMTModel, MarianTokenizer
from PIL import Image
import pinecone
import os
import requests
from io import BytesIO
from lxml import etree
from tqdm import tqdm
    
#2. Проверка доступного устройства
#Мы проверяем, доступна ли GPU для наших вычислений. Если нет, мы будем использовать CPU.

device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)
    
#3. Настройка модели перевода
#Мы будем использовать модель MarianMT для перевода с русского на английский. Модель и её токенизатор загружаются и перемещаются на доступное устройство.

model_name = 'Helsinki-NLP/opus-mt-ru-en'
tr_tokenizer = MarianTokenizer.from_pretrained(model_name)
tr_model = MarianMTModel.from_pretrained(model_name).to(device)
    
#4. Инициализация Pinecone
#Мы инициализируем Pinecone с помощью вашего API-ключа. Затем создаем новый индекс в Pinecone, если он еще не существует.

api_key = "YOUR_API_KEY"
pinecone.init(api_key, environment="us-central1-gcp")
index_name = "clip-embeddings"
if index_name not in pinecone.list_indexes():
    pinecone.create_index(
      index_name,
      dimension=1536,
      metric="cosine",
      pod_type="s1"
    )
index = pinecone.Index(index_name)
    
#5. Работа с XML лентой
#Мы загружаем XML ленту, извлекаем нужные поля и загружаем изображения.

from open_clip import tokenizer
url = "https://www.stolplit.ru/yml/mb_moscow.xml"
response = requests.get(url)
xml_content = response.content
root = etree.fromstring(xml_content)
print('Загружено')
    
#6. Обработка и загрузка изображений в Pinecone
#Мы обрабатываем изображения, вычисляем векторные представления и загружаем их в Pinecone. Это происходит в партиях для повышения эффективности.

offers = root.xpath("//offers/offer")
BATCH_SIZE = 100
offers = offers[59500:]
    
#7. Вычисление эмбеддингов и загрузка в Pinecone
#Мы предобрабатываем изображения, вычисляем эмбеддинги и объединяем их. Затем мы загружаем эмбеддинги в Pinecone.

with torch.no_grad():
    image_embeddings = im_model.encode_image(preprocessed_images)
    text_embeddings = im_model.encode_text(text_inputs)
combined_embeddings = torch.cat((image_embeddings, text_embeddings), dim=1)
    
#8. Загрузка в Pinecone
#Мы готовим данные и загружаем их в Pinecone.


upserts = []
for name, name_en, autoname_en, embedding, image_url, image_descr, product_url in zip(image_names, image_names_en, image_autocapts_en, combined_embeddings, image_urls, image_descriptions, product_urls):
    idx_s = str(hash(image_url))
    image_descr = image_descr if image_descr is not None else ""
    upserts.append({
        "id": idx_s,
        "values": embedding.tolist(),
        "metadata": {"name_ru": name, "name_en": name_en, "autoname_en":autoname_en, "image_url": image_url, 'image_descr': image_descr, 'product_url':product_url},
    })
index.upsert(upserts)