solution.py

import os
import json
import re
import pickle
from typing import Optional

import pandas as pd
import numpy as np
import scipy.sparse as sp
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import f1_score
from catboost import CatBoostClassifier
from openai import OpenAI

# --- CONFIG ---
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "")
OPENROUTER_BASE_URL = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
MODEL = os.getenv("OPENROUTER_MODEL", "openai/gpt-4o-mini")
MC_DICT_FILE = "rnc_mic_key_phrases.csv"
DATASET_FILE = "rnc_dataset_markup.json"
MODEL_PATH = "model/split_detector.pkl"

def load_micro_categories(filepath=MC_DICT_FILE):
    df = pd.read_csv(filepath)
    microcats = {}
    for _, row in df.iterrows():
        mc_id = int(row["mcId"])
        
        # Получаем строку и превращаем её в список, разделяя по ';'
        raw_phrases = str(row.get("keyPhrases", ""))
        if raw_phrases and raw_phrases.lower() != 'nan':
            phrases_list = [p.strip() for p in raw_phrases.split(';') if p.strip()]
        else:
            phrases_list = []
            
        microcats[mc_id] = {
            "mcId": mc_id, 
            "mcTitle": row["mcTitle"],
            "keyPhrases": phrases_list
        }
    return microcats

class Pipeline:
    def __init__(self, mc_dict):
        self.mc_dict = mc_dict
        self.client: Optional[OpenAI] = (
            OpenAI(api_key=OPENROUTER_API_KEY, base_url=OPENROUTER_BASE_URL)
            if OPENROUTER_API_KEY
            else None
        )
        
        # --- ЗОЛОТОЙ ДЕТЕКТОР ---
        self.regex_map = {
            102: r'(?i)(?<![а-яёa-z])(сантех|разводка труб|замена труб|унитаз|смесител|сифон|раковин|ванн[а-яё]*|душев|канализ|водопровод|водоснабжен|бойлер|радиатор|инсталл|инстал|полотенцесушител|мойк|моек|полотенчик)[а-яёa-z]*(?![а-яёa-z])',
            103: r'(?i)(?<![а-яёa-z])(электр|проводк|розетк|щиток|штробл|автомат\b|выключател|светильник|люстр|освещени|светодиод)[а-яёa-z]*(?![а-яёa-z])',
            104: r'(?i)(?<![а-яёa-z])(натяжн|потолок|потолк|парящ|тканев)[а-яёa-z]*(?![а-яёa-z])',
            105: r'(?i)(?<![а-яёa-z])(плитк|плиточн|кафел|керамогранит|мозаик|затирк|фартук)[а-яёa-z]*(?![а-яёa-z])',
            106: r'(?i)(?<![а-яёa-z])(обои|обоев|обой|обоям|флизелин|винил|поклейк|оклейк|наклеиван)[а-яёa-z]*(?![а-яёa-z])',
            107: r'(?i)(?<![а-яёa-z])(маляр|покрас|окрас|грунтов|побелк)[а-яёa-z]*(?![а-яёa-z])',
            108: r'(?i)(?<![а-яёa-z])(штукатур|ротбанд|выравниван|маяк|правил|шпакл[её]в|шпатл[её]в|откос)[а-яёa-z]*(?![а-яёa-z])',
            109: r'(?i)(?<![а-яёa-z])(ламинат|линол[еи]ум|паркет|стяжк|наливн|кварцвинил|плинтус|ковролин|тепл.*?пол|напольн)[а-яёa-z]*(?![а-яёa-z])',
            110: r'(?i)(?<![а-яёa-z])(гипс[ао]картон|гкл|гвл|гипрок|короб|перегородк|арк[аиу])[а-яёa-z]*(?![а-яёa-z])',
            111: r'(?i)(?<![а-яёa-z])(демонтаж|снос|разборк|алмазн)[а-яёa-z]*(?![а-яёa-z])'
        }
        
        self._init_ml()

    def detect_categories(self, text, source_mc_id):
        detected =[]
        for mc_id, pattern in self.regex_map.items():
            if mc_id == source_mc_id: 
                continue
            if re.search(pattern, text):
                detected.append(mc_id)
        return sorted(detected)

    def _extract_features(self, text, det_list):
        text_l = text.lower().replace('ё', 'е')
        wc = len(text_l.split()) + 1
        has_otdelno = 1 if re.search(r'отдельн|частичн|по отдельности|мелкий ремонт|разные виды', text_l) else 0
        has_complex = 1 if re.search(r'под ключ|комплекс|подчистовую|от а до я|с нуля', text_l) else 0
        list_markers = len(re.findall(r'[-•✔✅🔨✔️\*]', text))
        density = len(det_list) / wc
        has_price = 1 if re.search(r'цена|прайс|руб|₽|от \d+', text_l) else 0
        return[wc, has_otdelno, has_complex, len(det_list), list_markers, density, has_price]

    def _init_ml(self):
        if os.path.exists(MODEL_PATH):
            with open(MODEL_PATH, "rb") as f:
                art = pickle.load(f)
                self.tfidf, self.scaler, self.clf, self.threshold = art["tfidf"], art["scaler"], art["clf"], art["threshold"]
            return

        print("⚠️ Модель не найдена. Обучаем CatBoost + Smart Regex Features...")
        with open(DATASET_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
        
        texts, y, meta, det_lists =[], [], [], []
        for item in data:
            txt = str(item["description"]).lower().replace('ё', 'е')
            det = self.detect_categories(txt, item.get("sourceMcId", 101))
            
            texts.append(txt)
            y.append(1 if item.get("shouldSplit") else 0)
            det_lists.append(det)
            meta.append(self._extract_features(txt, det))

        self.tfidf = TfidfVectorizer(max_features=5000, ngram_range=(1, 3), sublinear_tf=True)
        X_tfidf = self.tfidf.fit_transform(texts)
        self.scaler = StandardScaler()
        X_meta = self.scaler.fit_transform(meta)
        X_final = sp.hstack([X_tfidf, X_meta])

        # Обучаем CatBoost на всех данных
        self.clf = CatBoostClassifier(
            iterations=600, learning_rate=0.08, depth=6, 
            auto_class_weights='SqrtBalanced', verbose=0, random_seed=42
        )
        self.clf.fit(X_final, y)

        # Авто-подбор порога на обучающей выборке
        y_prob = self.clf.predict_proba(X_final)[:, 1]
        best_threshold, best_f1 = 0.5, 0
        for t in np.arange(0.3, 0.8, 0.02):
            preds = (y_prob >= t).astype(int)
            for i, det in enumerate(det_lists):
                if len(det) == 0: preds[i] = 0
            score = f1_score(y, preds)
            if score > best_f1:
                best_f1, best_threshold = score, t

        self.threshold = best_threshold
        print(f"✅ Модель обучена! Идеальный порог: {self.threshold:.2f}")

        os.makedirs("model", exist_ok=True)
        with open(MODEL_PATH, "wb") as f:
            pickle.dump({"tfidf": self.tfidf, "scaler": self.scaler, "clf": self.clf, "threshold": self.threshold}, f)

    def generate_draft_text(self, original_text, mc_title):
        if self.client is None:
            return self._fallback_draft_text(mc_title)

        prompt = f"Напиши черновик объявления Авито для услуги: '{mc_title}'. Используй стиль автора: {original_text[:500]}"
        try:
            resp = self.client.chat.completions.create(model=MODEL, messages=[{"role": "user", "content": prompt}], temperature=0.3, max_tokens=150)
            return resp.choices[0].message.content.strip().replace('"', '')
        except Exception:
            return self._fallback_draft_text(mc_title)

    @staticmethod
    def _fallback_draft_text(mc_title):
        return f"Профессиональные услуги: {mc_title}. Качественно и в срок. Звоните!"

    def process(self, item, generate_text=True):
        desc = item["description"]
        det_ids = self.detect_categories(desc, item.get("sourceMcId", 101))
        
        # ML Inference
        x_tf = self.tfidf.transform([desc.lower().replace('ё', 'е')])
        x_mt = self.scaler.transform([self._extract_features(desc, det_ids)])
        
        should_split = False
        if det_ids:
            # Получаем вероятность и сравниваем с нашим умным порогом
            prob = self.clf.predict_proba(sp.hstack([x_tf, x_mt]))[0, 1]
            should_split = bool(prob >= self.threshold)

        drafts =[]
        if should_split:
            for mid in det_ids:
                txt = self.generate_draft_text(desc, self.mc_dict[mid]["mcTitle"]) if generate_text else ""
                drafts.append({"mcId": mid, "mcTitle": self.mc_dict[mid]["mcTitle"], "text": txt})

        return {
            "itemId": item["itemId"],
            "detectedMcIds": det_ids,
            "shouldSplit": should_split,
            "splitMcIds": det_ids if should_split else [],  # <--- Добавили эту строку
            "drafts": drafts
        }

if __name__ == "__main__":
    # Быстрая проверка
    mc = load_micro_categories()
    p = Pipeline(mc)
    print("🚀 Pipeline готов к работе и загружен в память!")