from sentence_transformers import SentenceTransformer
from wdr_article_semantic_chunking_2.util.util import generate_combinations, get_path_for_plotting, get_path_for_csv, \
    get_path_for_plotting2

from wdr_article_semantic_chunking_2.computer.plotter import plot_sliding_window
from wdr_article_semantic_chunking_2.computer.slid_win import compute, segment_topics_window
from wdr_article_semantic_chunking_2.text_util.sw_writer import save_pair_to_csv
from wdr_article_semantic_chunking_2.text_util.text_extractor import extract_texts_and_write_to_file


def find_boundaries(
        model_name="paraphrase-multilingual-mpnet-base-v2",
        window_size=3,
        min_gap=3,
        sentences: list = None,
        debug=False,
        expected_boundaries: list = None,
        result_file_name=None
):
    assert sentences is not None
    model = SentenceTransformer(model_name)

    boundaries, scores, indices = segment_topics_window(sentences, window_size=window_size, min_gap=min_gap,
                                                        model=model)
    if debug:
        plot_sliding_window(
            indices=indices,
            scores=scores,
            expected_boundaries=expected_boundaries,
            file_name=get_path_for_plotting2(result_file_name),
            window_size=window_size
        )

        save_pair_to_csv(
            key=result_file_name,
            possible_breaks=boundaries,
            expected_boundary=expected_boundaries,
            source_count=None,
            file_path=get_path_for_csv(result_file_name)
        )
        return (
            boundaries,
            scores,
            indices,
        )
    else:
        return boundaries


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",
    ]
    file_name = f"merged_filtered_{1}.json"
    blocks, expected_boundary, source_count, _ = extract_texts_and_write_to_file(file_name, False)

    # for model in models:
    #     for window_and_gap_comb in generate_combinations():# we generate different combination ( in total 36) of window_size and min_gap
    r = find_boundaries(sentences=blocks, expected_boundaries=expected_boundary, result_file_name=file_name)
    print(r)
