import pandas as pd

from wdr_article_semantic_chunking_2.text_util.sw_writer import save_pair_to_csv, MATCH_PERCENTAGE, save_result_tocsv
from llama_index.core.node_parser import SemanticSplitterNodeParser, SemanticDoubleMergingSplitterNodeParser, \
    LanguageConfig
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from wdr_article_semantic_chunking_2.text_util.text_extractor import extract_texts_and_write_to_file
from llama_index.core import Document, Settings

from wdr_article_semantic_chunking_2.util.util import get_path_for_csv


def compute(
        model_name,
        initial_threshold=0.45,
        appending_threshold=0.35,
        merging_threshold=0.35,
        max_chunk_size=1200,
        merging_range=2,
):
    Settings.embed_model = HuggingFaceEmbedding(
        model_name=model_name
    )
    combination_name = f"model_{model_name}_initTresh_{initial_threshold}_apTresh_{appending_threshold}_mTresh_{merging_threshold}_maxChunk_{max_chunk_size}_mRange_{merging_range}"
    for i in range(0, 100):
        file_name = f"merged_filtered_{i}.json"
        print(file_name)
        blocks, expected_boundary, source_count, debug_dict = extract_texts_and_write_to_file(file_name, False)

        splitter = SemanticDoubleMergingSplitterNodeParser(
            language_config=LanguageConfig(language="german", spacy_model="de_core_news_md"),
            sentence_splitter=lambda text: blocks,

            # ---- thresholds ----
            initial_threshold=initial_threshold,
            appending_threshold=appending_threshold,
            merging_threshold=merging_threshold,

            # ---- merging ----
            max_chunk_size=max_chunk_size,
            merging_range=merging_range,
            merging_separator=" ",
        )
        doc = Document(text=" ".join(blocks))
        nodes = splitter.get_nodes_from_documents([doc])
        inverted = {v: k for k, v in debug_dict.items()}
        boundaries = get_approximate_boundary(nodes, debug_dict)

        save_pair_to_csv(
            key=file_name,
            possible_breaks=boundaries,
            expected_boundary=expected_boundary,
            source_count=source_count,
            file_path=get_path_for_csv(combination_name)
        )
    df = pd.read_csv(get_path_for_csv(combination_name), usecols=[MATCH_PERCENTAGE])
    save_result_tocsv(combination_name, df.mean().iloc[0])


def get_approximate_boundary(nodes, dict_for_debug):
    print(f"Total Chunks: {len(nodes)}")
    aproximate_index = []
    for i, node in enumerate(nodes, start=0):
        for t in range(0, max(dict_for_debug.keys())):
            debug = dict_for_debug[t].split()[:10]
            debug_sub1 = dict_for_debug.get(t - 1, "").split()[:10]
            debug_sub2 = dict_for_debug.get(t - 2, "").split()[:10]
            debug_add1 = dict_for_debug.get(t + 1, "").split()[:10]
            debug_add2 = dict_for_debug.get(t + 2, "").split()[:10]

            if len(debug) > 0 and " ".join(node.get_content().split()[:len(debug)]) == " ".join((debug)):
                aproximate_index.append(t)
                break
            if len(debug_sub1) > 0 and " ".join(node.get_content().split()[:len(debug_sub1)]) == " ".join((debug_sub1)):
                aproximate_index.append(t - 1)
                break
            if len(debug_sub2) > 0 and " ".join(node.get_content().split()[:len(debug_sub2)]) == " ".join((debug_sub2)):
                aproximate_index.append(t - 2)
                break

            if len(debug_add1) > 0 and " ".join(node.get_content().split()[:len(debug_add1)]) == " ".join((debug_add1)):
                aproximate_index.append(t + 1)
                break
            if len(debug_add2) > 0 and " ".join(node.get_content().split()[:len(debug_add2)]) == " ".join((debug_add2)):
                aproximate_index.append(t + 2)
                break
    return aproximate_index


if __name__ == "__main__":
    models = [
        # "all-MiniLM-L12-v2",
        # "all-MiniLM-L6-v2",
        # "all-distilroberta-v1",
        # "all-mpnet-base-v2",
        # "distiluse-base-multilingual-cased-v1",
        # "distiluse-base-multilingual-cased-v2",
        # "multi-qa-MiniLM-L6-cos-v1",
        # "multi-qa-distilbert-cos-v1",
        # "multi-qa-mpnet-base-dot-v1",
        # "paraphrase-MiniLM-L3-v2",
        # "paraphrase-albert-small-v2",
        "paraphrase-multilingual-MiniLM-L12-v2",
        # "paraphrase-multilingual-mpnet-base-v2",
    ]
    for model in models:
        # combination_name = f"model_{model}_buf_{0}_t_{80}"
        # df = pd.read_csv(get_path_for_csv(combination_name), usecols=[MATCH_PERCENTAGE])
        # save_result_tocsv(combination_name, df.mean().iloc[0])
        compute(
            initial_threshold=0.9,
            appending_threshold=0.20,
            merging_threshold=0.40,
            max_chunk_size=1500,
            merging_range=2,
            model_name=model)
