import re
import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity


class Segmenter:
    def __init__(self):
        print("Loading models :")

        self.model_bge = SentenceTransformer("C:/Models/bge-m3", model_kwargs={"use_safetensors": False})
        self.model_e5 = SentenceTransformer("C:/Models/multilingual-e5-large")
        self.model_labse = SentenceTransformer("C:/Models/LaBSE_Local")
        self.model_parsroberta = SentenceTransformer("C:/Models/ParsRoberta")

        print("4 models loaded successfully")

    def get_similarities(self, text1, text2):
        """محاسبه میانگین شباهت کسینوسی از ۴ مدل برای دو متن داده شده"""
        emb_bge = self.model_bge.encode([text1, text2])
        emb_e5 = self.model_e5.encode([text1, text2])
        emb_labse = self.model_labse.encode([text1, text2])
        emb_parsroberta = self.model_parsroberta.encode([text1, text2])

        sim_bge = cosine_similarity([emb_bge[0]], [emb_bge[1]])[0][0]
        sim_e5 = cosine_similarity([emb_e5[0]], [emb_e5[1]])[0][0]
        sim_labse = cosine_similarity([emb_labse[0]], [emb_labse[1]])[0][0]
        sim_parsroberta = cosine_similarity([emb_parsroberta[0]], [emb_parsroberta[1]])[0][0]

        # میانگین‌گیری از هر ۴ مدل
        return np.mean([sim_bge, sim_e5, sim_labse, sim_parsroberta])

    def calculate_depth_scores(self, similarities):
        """محاسبه عمق دره‌ها (Gap Scores) با استفاده از قله‌های چپ و راست"""
        depth_scores = []
        n = len(similarities)

        for i in range(n):
            left_peak = similarities[i]
            for j in range(i, -1, -1):
                if similarities[j] >= left_peak:
                    left_peak = similarities[j]
                else:
                    break

            right_peak = similarities[i]
            for j in range(i, n):
                if similarities[j] >= right_peak:
                    right_peak = similarities[j]
                else:
                    break

            depth = (left_peak - similarities[i]) + (right_peak - similarities[i])
            depth_scores.append(depth)

        return depth_scores

    def extract_sentences(self, text):
        """جداسازی متن اولیه به جملات"""
        text = text.replace('\n', ' ')
        sentences = [s.strip() for s in re.split(r'(?<=[.!?؟]) +', text) if s.strip()]
        return sentences

    def segment_items(self, items, force_count=None, std_multiplier=0.5, min_depth=0.02):
        """
        دریافت جملات یا واحدهای معنایی و کلاسترینگ آن‌ها.
        اگر پارامتر force_count ارسال شود، الگوریتم موظف است دقیقاً همان تعداد سگمنت تولید کند.
        """
        n_items = len(items)
        if n_items < 2:
            return [items]

        similarities = []
        for i in range(n_items - 1):
            sim = self.get_similarities(items[i], items[i + 1])
            similarities.append(sim)

        depth_scores = self.calculate_depth_scores(similarities)

        # -------------------------------------------------------------
        # حالت اجباری: رساندن تعداد سگمنت‌ها به عدد دقیق (مثلا 7)
        # -------------------------------------------------------------
        if force_count and 1 < force_count < n_items:
            # پیدا کردن عمیق‌ترین دره‌ها برای ایجاد دقیقا force_count - 1 برش
            top_cut_indices = np.argsort(depth_scores)[-(force_count - 1):]
            top_cut_indices = sorted(top_cut_indices)

            grouped_segments = []
            start = 0
            for idx in top_cut_indices:
                grouped_segments.append(items[start:idx + 1])
                start = idx + 1
            grouped_segments.append(items[start:])
            return grouped_segments

        # -------------------------------------------------------------
        # حالت عادی داینامیک: در سطوح پایین‌تر که هنوز به تعداد هدف نرسیده‌ایم
        # -------------------------------------------------------------
        mean_depth = np.mean(depth_scores)
        std_depth = np.std(depth_scores)
        cutoff_threshold = mean_depth + (std_multiplier * std_depth)

        grouped_segments = []
        current_segment = [items[0]]

        for i in range(len(similarities)):
            if depth_scores[i] >= cutoff_threshold and depth_scores[i] > min_depth:
                grouped_segments.append(current_segment)
                current_segment = [items[i + 1]]
            else:
                current_segment.append(items[i + 1])

        if current_segment:
            grouped_segments.append(current_segment)

        return grouped_segments