"""Test of effects of sentence length

This script measures the length of sentences in a data set, gives statistics about the length (average, standard
deviation, distribution) and most importantly creates three sub-datasets according to the length of a data point's
sentences. This is done to measure the model's performance on sentences of different lengths.
"""
from src.SBERT_Model import SBERT_Model
from STSB_Dataset import STSB_Dataset
import math
import pandas


def add_length_column(dataset) -> pandas.DataFrame:
    """Creates new column containing length of sentences

    A new column is appended to the given dataframe. It consists of the length of both sentences of
    this row combined.
    """
    sentence_1_lengths = dataset['sentence1'].apply(lambda sentence: len(sentence.split(" ")))
    sentence_2_lengths = dataset['sentence2'].apply(lambda sentence: len(sentence.split(" ")))
    dataset['total length'] = sentence_1_lengths
    dataset['total length'] += sentence_2_lengths
    return dataset


def get_length_distribution(dataset) -> dict:
    """Counts how often each sentence length appears

    For the given dataframe, first combines both 'sentence' columns, then counts length of this column for each row.
    Returns counts of every found length as a dict.
    """
    dataset['sentences combined'] = dataset['sentence1'] + " " + dataset['sentence2']
    sentences = dataset['sentences combined'].to_numpy()
    lengths = dict()
    for sentence in sentences:
        sentence_length = len(sentence.split(" "))
        if sentence_length not in lengths:
            lengths[sentence_length] = 0
        lengths[sentence_length] += 1
    return lengths


if __name__ == "__main__":
    complete_stsb_dataset = STSB_Dataset('../data/stsbenchmark/sts-test.csv').as_dataframe  # adjust used dataset here
    dataset_with_lengths = add_length_column(complete_stsb_dataset)

    average_length = dataset_with_lengths['total length'].mean()
    standard_deviation = dataset_with_lengths['total length'].std()
    print(f"Average length: {average_length}, standard deviation: {standard_deviation}")  # first print precise values
    average_length = math.ceil(average_length)  # then use next highest int to get usable values
    standard_deviation = math.ceil(standard_deviation)

    for length, occurrence in sorted(get_length_distribution(dataset_with_lengths).items()):
        print(f"Length: {length} | Occurrence: {occurrence}")

    # create data sets corresponding to sentence length
    short_sentences = dataset_with_lengths.loc[
        dataset_with_lengths['total length'] <= (average_length - standard_deviation)]
    medium_length_sentences = dataset_with_lengths.loc[
        (dataset_with_lengths['total length'] > (average_length - standard_deviation)) &
        (dataset_with_lengths['total length'] < (average_length + standard_deviation))]
    long_sentences = dataset_with_lengths.loc[
        dataset_with_lengths['total length'] >= (average_length + standard_deviation)]

    datasets = [short_sentences, medium_length_sentences, long_sentences]

    for d in datasets:
        ft_model = SBERT_Model("Fine-tuned Model", '../models/stsb-bert-large/', d)
        ft_model.print_statistics()
