from pyvis.network import Network
import pandas as pd
import pymorphy3
import re
import gradio as gr

# Инициализация pymorphy3 (лемматизатор)
morph = pymorphy3.MorphAnalyzer()

# Функция токенизации и лемматизации
def tokenize_and_lemmatize(text):
    text = re.sub(r'[^\w\s]', '', text.lower())  # Удаляем пунктуацию, приводим к нижнему регистру
    words = text.split()
    return [morph.parse(word)[0].normal_form for word in words]

# Словарь ключевых слов для каждого сегмента (каждое слово в начальной форме)
personalization_keywords = {
    "Поколение X": ["комиссия", "визит", "снижение", "ставка", "бесплатно", "экономия"],
    "Поколение Y": ["онлайн", "цифровой", "бонус", "лимит", "qr", "sberpay"],
    "Поколение Z": ["быстрота", "мгновенно", "минута", "оперативно", "решение"],
    "Пол Женский": ["комфорт", "удобство", "забота", "легкость", "терминал"],
    "ОПФ ИП": ["эффективность", "комиссия", "снижение", "ставка", "онлайн", "быстрота", "оптимизация"],
    "ОПФ ООО": ["удобство", "комиссия", "открытие", "онлайн", "легкость", "автоматизация"],
    "Психотип Конструктор": ["оптимизация", "снижение", "ставка", "комиссия", "льготный", "период", "выгодно", "настройка"],
    "Пол Мужской": ["динамичность", "быстрота", "лимит", "решение", "активность", "оптимизация"],
    "Стадия бизнеса Новичок": ["доступность", "простота", "низкий", "порог", "комиссия", "легкость"],
    "Стадия бизнеса Профи": ["профессионализм", "комиссия", "снижение", "ставка", "оптимизация", "эффективность", "быстрота"],
    "Психотип Рефлектор": ["премиальность", "эксклюзив", "бизнес", "зал", "акция", "привилегия", "статус"],
    "Психотип Центрист": ["универсальность", "стандарт", "комиссия", "бесплатно", "надежность"],
    "Стадия бизнеса Эксперт": ["максимизация", "высокий", "лимит", "снижение", "ставка", "комиссия", "выгода", "оптимизация"]
}

# Глобальная переменная для хранения DataFrame
df = None

def load_excel(file):
    global df
    if file is None:
        return [], "Файл не загружен. Загрузите Excel-файл."
    try:
        df = pd.read_excel(file.name, usecols=["Продукт", "Преимущество"])
        unique_products = df["Продукт"].unique().tolist()
        return unique_products, "Файл успешно загружен!"
    except Exception as e:
        return [], f"Ошибка при чтении файла: {str(e)}"

def classify_advantage(text, keywords_dict):
    lemmas = tokenize_and_lemmatize(text)
    category_matches = {}

    for category, keywords in keywords_dict.items():
        matches = set(lemmas) & set(keywords)
        if matches:
            category_matches[category] = {
                'count': len(matches),
                'matched_lemmas': matches
            }

    sorted_matches = sorted(category_matches.items(), key=lambda x: x[1]['count'], reverse=True)
    return sorted_matches

def analyze(product):
    global df
    if df is None:
        return "Сначала загрузите файл.", None

    if not product:
        return "Пожалуйста, выберите продукт.", None

    product_advantages = df[df["Продукт"] == product]["Преимущество"]

    if product_advantages.empty:
        return "Для выбранного продукта не найдено преимуществ.", None

    graph_html = create_category_graph(product, product_advantages, personalization_keywords)

    results = []
    for advantage in product_advantages:
        matches = classify_advantage(advantage, personalization_keywords)
        advantage_text = f"**Преимущество**: {advantage}\n\n"
        advantage_text += f"**Леммы**: {tokenize_and_lemmatize(advantage)}\n\n"
        advantage_text += "**Совпадающие категории:**\n"

        if matches:
            for category, data in matches:
                matched_lemmas_str = ", ".join(sorted(data['matched_lemmas']))
                advantage_text += f"- {category}: {data['count']} совпадений (леммы: {matched_lemmas_str})\n"
        else:
            advantage_text += "- Нет совпадений.\n"
        advantage_text += "\n---\n"
        results.append(advantage_text)

    return "\n".join(results), graph_html

def create_category_graph(product, advantages, personalization_keywords):
    net = Network(height="500px", width="100%", directed=True, cdn_resources='in_line')

    net.add_node(product, label=product, color="lightblue", size=30)

    for advantage in advantages:
        net.add_node(advantage, label=advantage, color="orange", size=20)
        net.add_edge(product, advantage)

        matches = classify_advantage(advantage, personalization_keywords)
        for category, data in matches:
            net.add_node(category, label=category, color="green", size=15)
            net.add_edge(advantage, category)

    html = net.generate_html(notebook=False)
    html_escaped = html.replace('"', '"').replace("'", "'")

    iframe_html = f"""
    <iframe 
        width="100%" 
        height="600" 
        frameborder="0" 
        srcdoc="{html_escaped}">
    </iframe>
    """

    print("Generated HTML:", iframe_html[:500])  # Выводим первые 500 символов для отладки

    return iframe_html

with gr.Blocks() as demo:
    gr.Markdown("## Классификация преимуществ по признакам персонализации")
    gr.Markdown("**Шаг 1:** Загрузите Excel-файл с двумя столбцами: 'Продукт' и 'Преимущество'.")

    file_input = gr.File(label="Загрузите Excel-файл", file_types=[".xlsx"])
    load_button = gr.Button("Загрузить файл")
    load_status = gr.Markdown("")

    gr.Markdown("**Шаг 2:** Выберите продукт из списка.")
    product_dropdown = gr.Dropdown(choices=[], label="Продукты", value=None)
    analyze_button = gr.Button("Анализировать")

    output_text = gr.Markdown("")
    output_graph = gr.HTML(label="Визуализация графа")

    def on_file_upload(file):
        unique_products, status_message = load_excel(file)
        return gr.update(choices=unique_products), status_message

    load_button.click(
        fn=on_file_upload,
        inputs=file_input,
        outputs=[product_dropdown, load_status]
    )

    analyze_button.click(
        fn=analyze,
        inputs=product_dropdown,
        outputs=[output_text, output_graph]
    )

if __name__ == "__main__":
    demo.launch(debug=True)