TgBotMedical

# -*- coding: utf-8 -*-
"""
Полноценная реализация чат-бота для Telegram.
Бот обрабатывает входящие сообщения следующим образом:
1. Получает сообщение от пользователя.
2. С помощью NLTK:
   - Токенизирует текст,
   - Удаляет стоп-слова,
   - Оставляет только слова (без знаков препинания и чисел).
3. Приводит слова к нормальной форме (лемматизация) с помощью spaCy.
4. Дополнительно для симптомов пытается преобразовать глагол в существительное (например, "тошнить" -> "тошнота") с помощью pymorphy2.
5. Загружает таблицу "Лекарства.xlsx" с данными
   (столбцы: раздел, подраздел, наименование, Фармакологическое действие, Показания к применению, Побочные эффекты).
6. Производится поиск по всем столбцам.
7. Если по столбцу "наименование" находятся совпадения, бот выводит красиво отформатированную справку с эмодзи.
8. Если препарата не найдено, но совпадения обнаружены в столбцах "Показания к применению" или "Побочных эффектов",
   выводится сообщение с рекомендациями.
9. Если совпадений нет – выводится сообщение "Информация по запросу не найдена."
10. Бот также обрабатывает команды /start, /help и "Выход" для завершения работы.
"""

import os
import re
import telebot
import pandas as pd
import nltk
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
import spacy
import pymorphy2

# Загружаем необходимые ресурсы NLTK
nltk.download('punkt')
nltk.download('stopwords')

# Загружаем spaCy модель для русского языка
nlp = spacy.load("ru_core_news_md")

# Инициализируем pymorphy2 для морфологического анализа
morph = pymorphy2.MorphAnalyzer()

# Загружаем Excel-файл с лекарствами.
try:
    df = pd.read_excel("Лекарства.xlsx")
except Exception as e:
    print("Ошибка при загрузке файла Excel:", e)
    df = None

# Получаем токен из переменной окружения
TOKEN = os.getenv('TELEGRAM_BOT_TOKEN')
if TOKEN is None:
    raise ValueError("Переменная окружения TELEGRAM_BOT_TOKEN не установлена!")

bot = telebot.TeleBot(TOKEN)


def process_message(text):
    """
    Обрабатывает входящий текст:
      - Приводит к нижнему регистру.
      - Токенизирует, удаляет стоп-слова, оставляя только буквы.
      - Лемматизирует текст с помощью spaCy.
      - Для каждого токена, если он является глаголом, пытается найти производную существительную форму.
      - Применяет обработку синонимов (например, "насморк" -> "ринит")
      - Обрабатывает симптомы: заменяет "болеть" на "боль".
      - Если встречаются и "боль" и локализационные слова (например, "живот", "рука", "кости"),
        локализационные токены сохраняются отдельно и удаляются из итогового списка для поиска.
    Возвращает кортеж: (список лемм для поиска, список найденных локаций).

    Изменения: строки примерно 40-70.
    """
    text = text.lower()
    tokens = word_tokenize(text, language="russian")
    # Фильтруем: оставляем только буквы, исключаем стоп-слова и токены длиной меньше 3 символов
    filtered_tokens = [token for token in tokens
                       if token.isalpha() and token not in stopwords.words("russian") and len(token) > 2]
    doc = nlp(" ".join(filtered_tokens))
    lemmas = []
    for token in doc:
        lemma = token.lemma_
        # Преобразование глаголов в существительные (например, "тошнить" -> "тошнота", "рвать" -> "рвота")
        p = morph.parse(lemma)[0]
        if 'VERB' in p.tag and lemma in ['тошнить', 'рвать']:
            mapping = {'тошнить': 'тошнота', 'рвать': 'рвота'}
            lemma = mapping.get(lemma, lemma)
        # Обработка просторечных выражений (например, "насморк" -> "ринит")
        colloquial_mapping = {
            'насморк': 'ринит'
        }
        if lemma in colloquial_mapping:
            lemma = colloquial_mapping[lemma]
        # Обработка симптомов: замена "болеть" на "боль"
        symptom_mapping = {
            'болеть': 'боль'
        }
        if lemma in symptom_mapping:
            lemma = symptom_mapping[lemma]
        lemmas.append(lemma)
    # Определяем локализационные ключевые слова
    location_candidates = ['живот', 'рука', 'кости']
    locations_found = [token for token in lemmas if token in location_candidates]
    # Если встречаются и "боль" и локализация, удаляем локализационные токены из списка для поиска
    if "боль" in lemmas and locations_found:
        for loc in locations_found:
            while loc in lemmas:
                lemmas.remove(loc)
    return lemmas, locations_found


def search_tokens_in_table(tokens, dataframe):
    """
    Ищет совпадения для каждого токена по всем столбцам.
    Возвращает два словаря:
      col_matches: {имя_столбца: количество совпадений}
      col_matches_rows: {имя_столбца: список индексов строк, где найдены совпадения}
    """
    col_matches = {col: 0 for col in dataframe.columns}
    col_matches_rows = {col: [] for col in dataframe.columns}
    for token in tokens:
        # Создаем шаблон для точного совпадения по слову с учетом границ слова
        pattern = re.compile(r'\b' + re.escape(token) + r'\b')
        for col in dataframe.columns:
            for idx, value in dataframe[col].items():
                if pd.notnull(value) and pattern.search(str(value).lower()):
                    col_matches[col] += 1
                    col_matches_rows[col].append(idx)
    return col_matches, col_matches_rows


def analyze_results_all(col_matches_rows, dataframe, original_query, locations_found):
    """
    Анализирует результаты поиска.

    Сценарий 1: Если по столбцу "наименование" найдены совпадения, формирует отформатированную справку:
       💊 *Название:* <наименование>
          📁 *Раздел:* <раздел>
          📂 *Подраздел:* <подраздел>
          🔹 *Фармакологическое действие:* <...>
          🔹 *Показания к применению:* <...>
          🔹 *Побочные эффекты:* <...>

    Сценарий 2: Если препарата не найдено, но есть совпадения в "Показаниях к применению" или "Побочных эффектах",
    формируется сообщение вида:
       "Возможно, у вас наблюдаются симптомы: <original_query>
        ⚠️ Это может быть побочным эффектом препарата: <перечень препаратов>
        💉 При таких симптомах показания по лечению: <перечень препаратов>"

    Если совпадений нет, возвращается сообщение "Информация по запросу не найдена."

    Дополнительно, если обнаружены локализации, связанные с травматологией (например, "рука" или "кости"),
    к сообщению добавляется рекомендация обратиться к травматологу.

    Изменения: строка примерно 110.
    """
    # Сценарий 1: Поиск по "наименование"
    if col_matches_rows.get("наименование") and len(col_matches_rows["наименование"]) > 0:
        matched_indices = set(col_matches_rows["наименование"])
        result_lines = []
        for idx in matched_indices:
            row = dataframe.loc[idx]
            med_name = row.get("наименование", "Без наименования")
            section = row.get("раздел", "нет данных")
            subsection = row.get("подраздел", "нет данных")
            pharm_action = row.get("Фармакологическое действие", "нет данных")
            indications = row.get("Показания к применению", "нет данных")
            side_effects = row.get("Побочные эффекты", "нет данных")
            line = (f"💊 *Название:* {med_name}\n"
                    f"   📁 *Раздел:* {section}\n"
                    f"   📂 *Подраздел:* {subsection}\n"
                    f"   🔹 *Фармакологическое действие:* {pharm_action}\n"
                    f"   🔹 *Показания к применению:* {indications}\n"
                    f"   🔹 *Побочные эффекты:* {side_effects}")
            result_lines.append(line)
        message_text = "Возможно, вы искали:\n\n" + "\n\n".join(result_lines)
    else:
        # Сценарий 2: Поиск по симптомам в "Показаниях к применению" и "Побочных эффектах"
        side_effects_indices = set(col_matches_rows.get("Побочные эффекты", []))
        indications_indices = set(col_matches_rows.get("Показания к применению", []))
        messages = []
        if side_effects_indices:
            meds_side = []
            for idx in side_effects_indices:
                row = dataframe.loc[idx]
                med_name = row.get("наименование", "Неизвестно")
                meds_side.append(med_name)
            meds_side = list(set(meds_side))
            messages.append("⚠️ Это может быть побочным эффектом препарата: " + ", ".join(meds_side))
        if indications_indices:
            meds_ind = []
            for idx in indications_indices:
                row = dataframe.loc[idx]
                med_name = row.get("наименование", "Неизвестно")
                meds_ind.append(med_name)
            meds_ind = list(set(meds_ind))
            messages.append("💉 При таких симптомах показания по лечению: " + ", ".join(meds_ind))
        if messages:
            message_text = (f"Возможно, у вас наблюдаются симптомы: *{original_query}*.\n" +
                            "\n".join(messages))
        else:
            message_text = "Информация по запросу не найдена."
    # Если в найденных локациях обнаружены слова, связанные с травматологией, добавляем рекомендацию
    if any(loc in ['рука', 'кости'] for loc in locations_found):
        message_text += "\n\nРекомендуем также обратиться к травматологу для уточнения диагноза."
    return message_text


@bot.message_handler(commands=['start', 'help'])
def send_welcome(message):
    bot.reply_to(message, 'Привет! Чем могу помочь? Напиши "Выход" для завершения работы.')


@bot.message_handler(content_types=['text'])
def get_text_messages(message):
    if message.text.strip().lower() == "выход":
        bot.send_message(message.from_user.id, "Бот завершает работу. До свидания!")
        exit_program()
        return
    original_query = message.text.strip()
    # Получаем леммы и найденные локации
    lemmas, locations_found = process_message(message.text)
    print("Обработанные токены:", lemmas, "Найденные локации:", locations_found)
    # Если запрос содержит "боль", но не указана локализация, просим уточнить
    if "боль" in lemmas and not locations_found:
        bot.send_message(message.from_user.id,
                         "Пожалуйста, уточните локализацию боли (например, 'боль в животе' или 'боль в руке').")
        return
    if df is None:
        bot.send_message(message.from_user.id, "Ошибка: не удалось загрузить таблицу с данными.")
        return
    col_matches, col_matches_rows = search_tokens_in_table(lemmas, df)
    print("Совпадения по столбцам:", col_matches)
    result = analyze_results_all(col_matches_rows, df, original_query, locations_found)
    bot.send_message(message.from_user.id, result, parse_mode="Markdown")


def exit_program():
    print("Завершение работы программы...")
    bot.stop_polling()
    quit()


bot.polling(none_stop=True, interval=0)