# 正则表达式
import re

from sentence_transformers import SentenceTransformer

import numpy as np

# 嵌入模型
model = SentenceTransformer("all-MiniLM-L6-v2")


class Semantic_splitter:
    def __init__(self, window_size, threshold):
        # 设置每个窗口的大小
        self.window_size = window_size
        # 设置相邻窗口的相似度的阈值
        self.threshold = threshold

    def create_documents(self, text):
        # 使用正则表达式对文本进行分割
        sentences = re.split(r"(。|！|？|\!|\?|\.|\n)", text)
        # print(sentences)
        # 初始化句子列表
        sents = []
        for i in range(0, len(sentences) - 1, 2):
            s = sentences[i].strip() + sentences[i + 1].strip()
            if s.strip():
                sents.append(s)

        print(sents)

        # 开始使用滑动窗口依据window_size进行切块
        # 初始化分块列表
        docs = []
        # 起始位置
        start = 0
        while start < len(sents):
            # 结束位置
            end = min(start + self.window_size, len(sents))
            window = sents[start:end]
            docs.append("".join(window))
            start = end
        print(f"{len(docs)}个块")  # 分为了3块，1句和2句,3句和4句,5句
        print(docs)
        # 计算每个块的向量值
        embeddings = model.encode(docs)
        # 初始化分割点列表，起点为0
        split_points = [0]
        print(embeddings)
        # 计算余弦相似度
        for i in range(1, len(docs)):
            sim = np.dot(
                embeddings[i - 1],
                embeddings[i]
                / (np.linalg.norm(embeddings[i - 1]) * np.linalg.norm(embeddings[i])),
            )
            if sim < self.threshold:
                print(f"相似度低于阈值{self.threshold},在位置{i}添加分割点")
                split_points.append(i)
        print(split_points)

        # 初始化最终分块列表
        result = []
        for i in range(len(split_points)):
            start = split_points[i]
            end = split_points[i + 1] if i + 1 < len(split_points) else len(docs)
            chunk = "".join(docs[start:end])
            if chunk.strip():
                result.append(chunk)

        return result


long_text = """今天天气晴朗，适合去公园散步。

量子力学中的叠加态是描述粒子同时处于多个状态的数学工具。

Windows命令行中复制文件可以使用copy命令。

大熊猫主要以竹子为食，是中国的国宝。

欧拉公式被誉为“最美的数学公式”。"""

semantic_spitter = Semantic_splitter(window_size=2, threshold=0.85)
documents = semantic_spitter.create_documents(long_text)
print(f"总共分割为{len(documents)}个块")

for i, doc in enumerate(documents, 1):
    print(f"---{i}个块")
    print(doc)