"""Unsupervised SimCSE on the FW make/model strings.

This is the simcse_script.ipynb pipeline pointed at our own data instead of
wiki1m/stsb. SimCSE's positive pair is the *same* sentence encoded twice -- the
two views differ only through dropout -- and MultipleNegativesRankingLoss
supplies the negatives from the rest of the batch.

Run prepare_data.py first.

  python simcse/train_simcse.py                      # defaults
  python simcse/train_simcse.py --epochs 10 --batch-size 30
"""

import argparse
import json
import logging
from datetime import datetime
from pathlib import Path
from xml.parsers.expat import model
from datasets import Dataset
from sentence_transformers import SentenceTransformer
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
from sentence_transformers.losses import MultipleNegativesRankingLoss
from sentence_transformers.models import Pooling, Transformer
from sentence_transformers.similarity_functions import SimilarityFunction
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments

HERE = Path(__file__).resolve().parent
DATA = HERE / "data"

logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO)


def read_jsonl(path: Path) -> list[dict]:
    with path.open() as f:
        return [json.loads(line) for line in f]


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--model-name", default="distilbert/distilbert-base-uncased")
    p.add_argument("--batch-size", type=int, default=30)
    p.add_argument("--epochs", type=int, default=10)
    # make/model strings are a handful of tokens; 32 is already generous
    p.add_argument("--max-seq-length", type=int, default=32)
    p.add_argument("--learning-rate", type=float, default=3e-5)
    p.add_argument("--output-dir", default=None)
    return p.parse_args()


def main():
    args = parse_args()
    out = Path(args.output_dir) if args.output_dir else HERE / "output" / (
        "simcse-make-model-{}-bs{}-{}".format(
            args.model_name.split("/")[-1],
            args.batch_size,
            datetime.now().strftime("%Y-%m-%d-%H-%M-%S"),
        )
    )

    # ---- model ----------------------------------------------------------
    word_embedding_model = Transformer(args.model_name, max_seq_length=args.max_seq_length)
    pooling_model = Pooling(word_embedding_model.get_sentence_embedding_dimension(), "mean")
    model = SentenceTransformer(modules=[word_embedding_model, pooling_model])

    # ---- data -----------------------------------------------------------
    train_rows = read_jsonl(DATA / "train.jsonl")
    train_dataset = Dataset.from_list(train_rows)

    eval_rows = read_jsonl(DATA / "eval.jsonl")

    evaluator = EmbeddingSimilarityEvaluator(
        sentences1=[r["sentence1"] for r in eval_rows],
        sentences2=[r["sentence2"] for r in eval_rows],
        scores=[r["score"] for r in eval_rows],
        main_similarity=SimilarityFunction.COSINE,
        name="make-model-dev",
    )

    logging.info("train sentences=%d  eval pairs=%d", len(train_dataset), len(eval_rows))

    loss = MultipleNegativesRankingLoss(model)

    logging.info("performance before training")
    before = evaluator(model)
    logging.info("%s", before)

    # ---- train ----------------------------------------------------------
    training_args = SentenceTransformerTrainingArguments(
        output_dir=str(out),
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        warmup_ratio=0.1,
        learning_rate=args.learning_rate,
        eval_strategy="steps",
        eval_steps=0.1,
        logging_steps=0.05,
        save_strategy="no",
        fp16=True,
        # identical strings in one batch become false negatives under MNRL
        batch_sampler="no_duplicates",
        report_to=[],
    )

    trainer = SentenceTransformerTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        evaluator=evaluator,
        loss=loss,
    )
    trainer.train()

    after = evaluator(model)
    logging.info("performance after training")
    logging.info("%s", after)

    final = out.parent / (out.name + "-final")
    model.save_pretrained(str(final))
    logging.info("saved -> %s", final)

    key = "make-model-dev_spearman_cosine"
    print("\n=== dev spearman (cosine) ===")
    print(f"  before : {before.get(key):.4f}")
    print(f"  after  : {after.get(key):.4f}")
    print(f"  model  : {final}")
   


if __name__ == "__main__":
    main()