from torch.utils.data import DataLoader
import math
from sentence_transformers import SentenceTransformer,  LoggingHandler, losses, models, util
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
from sentence_transformers.readers import InputExample
import logging
import sys
import csv


def read_sts(sts_dir, split):
    samples = []
    with open(f"{sts_dir}/sts-{split}_sk.csv", 'rt', encoding='utf8') as f:
        reader = csv.DictReader(f, delimiter='\t', quoting=csv.QUOTE_NONE)
        for row in reader:
            score = float(row['answer']) / 5.0
            inp_example = InputExample(texts=[row['question'], row['passage']], label=score)
            samples.append(inp_example)
    return samples


def main(model_name, output_dir, sts_dir):
    logging.basicConfig(format='%(asctime)s - %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S',
                        level=logging.INFO,
                        handlers=[LoggingHandler()])

    train_batch_size = 32
    num_epochs = 4

    word_embedding_model = models.Transformer(model_name)

    pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(),
                                   pooling_mode_mean_tokens=True,
                                   pooling_mode_cls_token=False,
                                   pooling_mode_max_tokens=False)

    model = SentenceTransformer(modules=[word_embedding_model, pooling_model])

    logging.info("Read STSbenchmark train dataset")

    train_samples = read_sts(sts_dir, 'train')
    dev_samples = read_sts(sts_dir, 'dev')
    test_samples = read_sts(sts_dir, 'test')

    train_dataloader = DataLoader(train_samples, shuffle=True, batch_size=train_batch_size)
    train_loss = losses.CosineSimilarityLoss(model=model)

    evaluator = EmbeddingSimilarityEvaluator.from_input_examples(dev_samples, name='sts-dev')

    warmup_steps = math.ceil(len(train_dataloader) * num_epochs * 0.1)
    logging.info("Warmup-steps: {}".format(warmup_steps))

    model.fit(train_objectives=[(train_dataloader, train_loss)],
              evaluator=evaluator,
              epochs=num_epochs,
              evaluation_steps=1000,
              warmup_steps=warmup_steps,
              output_path=output_dir)

    model = SentenceTransformer(output_dir)
    test_evaluator = EmbeddingSimilarityEvaluator.from_input_examples(test_samples, name='sts-test')
    test_evaluator(model, output_path=output_dir)


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2], sys.argv[3])

