import os
import sys
import torch
import logging
from datasets import load_dataset
from sentence_transformers import SentenceTransformer, losses
from sentence_transformers.evaluation import TripletEvaluator, SequentialEvaluator
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments, BatchSamplers
import glob

# --- CONFIG ---
MODEL_PATH = "/home/skiredj.abderrahman/khalil/sbert_training/epoch2/model/"
train_batch_size = 64
output_dir = "output/arabert_ms_marco"

# --- LOGGING ---
logging.basicConfig(
    format="%(asctime)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
    handlers=[logging.FileHandler("logs_ds36.txt")]
)

class Tee:
    def __init__(self, *files): self.files = files
    def write(self, obj):
        for f in self.files: f.write(obj); f.flush()
    def flush(self):
        for f in self.files: f.flush()
    def isatty(self): return False

sys.stdout = Tee(sys.stdout, open("logs_ds36.txt", "a"))

# --- LOAD DATA ---
train_dataset = load_dataset("csv", data_files="clean_dataset36_train.csv")["train"]
val_dataset = load_dataset("csv", data_files="clean_dataset36_val.csv")["train"]
print(f"Train size: {len(train_dataset)} | Val size: {len(val_dataset)}")

# --- MODEL ---
device = "cuda" if torch.cuda.is_available() else "cpu"
model = SentenceTransformer(MODEL_PATH, device=device)

# --- LOSS & EVALUATORS ---
matryoshka_dims = [768, 512, 256, 128, 64]
inner_train_loss = losses.MultipleNegativesRankingLoss(model=model)
train_loss = losses.MatryoshkaLoss(model, inner_train_loss, matryoshka_dims=matryoshka_dims)

evaluators = [
    TripletEvaluator(
        anchors=val_dataset["anchor"],
        positives=val_dataset["positive"],
        negatives=val_dataset["negative"],
        name=f"dev-{dim}",
        truncate_dim=dim,
    ) for dim in matryoshka_dims
]
dev_evaluator = SequentialEvaluator(evaluators, main_score_function=lambda scores: scores[0])

# --- TRAINING ARGS ---
args = SentenceTransformerTrainingArguments(
    output_dir=output_dir,
    num_train_epochs=2,
    per_device_train_batch_size=train_batch_size,
    gradient_accumulation_steps=2,
    bf16=True,
    learning_rate=1e-5,
    warmup_ratio=0.1,
    batch_sampler=BatchSamplers.NO_DUPLICATES,
    eval_strategy="steps",
    eval_steps=12000,
    save_strategy="steps",
    save_steps=12000,
    save_total_limit=2,
    logging_steps=200,
)

# --- RESUME IF CHECKPOINT EXISTS ---
existing_checkpoints = sorted(glob.glob(f"{output_dir}/checkpoint-*"))
resume_from = existing_checkpoints[-1] if existing_checkpoints else None
if resume_from:
    print(f"Resuming from: {resume_from}")
else:
    print("Starting fresh")

# --- TRAIN ---
trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    loss=train_loss,
    evaluator=dev_evaluator,
)
trainer.train(resume_from_checkpoint=resume_from)

# --- SAVE FINAL ---
final_output_dir = "/home/skiredj.abderrahman/khalil/sbert_training/output/final_ms_marco"
model.save(final_output_dir)
print(f"Model saved to {final_output_dir}")

# --- FINAL EVAL ---
test_evaluators = [
    TripletEvaluator(
        anchors=val_dataset["anchor"],
        positives=val_dataset["positive"],
        negatives=val_dataset["negative"],
        name=f"final-{dim}",
        truncate_dim=dim,
    ) for dim in matryoshka_dims
]
SequentialEvaluator(test_evaluators)(model)