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

from wdr_article_semantic_chunking_2.computer.plotter import plot_sliding_window
from wdr_article_semantic_chunking_2.text_util.sw_writer import save_pair_to_csv, MATCH_PERCENTAGE, save_result_tocsv
from wdr_article_semantic_chunking_2.text_util.text_extractor import extract_texts_and_write_to_file
from wdr_article_semantic_chunking_2.util.util import get_path_for_plotting, get_path_for_csv


def compute(
        window_size,
        min_gap,
        model_name):
    model = SentenceTransformer(model_name)
    combination_name = f"model_{model_name}_w_{window_size}_m_{min_gap}"
    for i in range(0, 100):
        file_name = f"merged_filtered_{i}.json"
        blocks, expected_boundary, source_count, _ = extract_texts_and_write_to_file(file_name, False)

        boundaries, scores, indices = segment_topics_window(blocks, window_size=window_size, min_gap=min_gap,
                                                            model=model)

        plot_sliding_window(
            indices=indices,
            scores=scores,
            expected_boundaries=expected_boundary,
            file_name=get_path_for_plotting(combination_name, file_name),
            window_size=window_size
        )

        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 segment_topics_window(
        blocks,
        window_size,
        min_gap,
        model
):
    embeddings = model.encode(blocks)

    scores = []
    indices = []

    for i in range(window_size, len(blocks) - window_size):
        left = embeddings[i - window_size:i]
        right = embeddings[i:i + window_size]

        left_mean = optimize_embddings(left)
        right_mean = optimize_embddings(right)

        sim = cosine_similarity(left_mean, right_mean)[0][0]
        scores.append(sim)
        indices.append(i)

    threshold = np.mean(scores) - 1.2 * np.std(scores)

    boundaries = []
    last = 0

    for idx, score in zip(indices, scores):
        if score < threshold and idx - last >= min_gap:
            boundaries.append(idx)
            last = idx

    return boundaries, scores, indices


def optimize_embddings(left):
    return np.mean(left, axis=0).reshape(1, -1)
