import logging
import sys
from datetime import datetime

import torch
from datasets import load_dataset, DatasetDict
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,
    MultiDatasetBatchSamplers,
)

# ─────────────────────────────────────────────────────────────
# Logging
# ─────────────────────────────────────────────────────────────
logging.basicConfig(
    format="%(asctime)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
    handlers=[logging.FileHandler("logs.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.txt", "a"))

# ─────────────────────────────────────────────────────────────
# Config
# ─────────────────────────────────────────────────────────────
train_batch_size = 32
model_name       = "/home/skiredj.abderrahman/khalil/sbert_training/third_training/output/arabert_20260326_1254/final"
model_nickname   = "arabert_arcd_squad"
timestamp        = datetime.now().strftime("%Y%m%d_%H%M")
output_dir       = f"output/{model_nickname}_{timestamp}"
device           = "cuda" if torch.cuda.is_available() else "cpu"
matryoshka_dims  = [768, 512, 256, 128, 64]

# ─────────────────────────────────────────────────────────────
# 1. Model
# ─────────────────────────────────────────────────────────────
model = SentenceTransformer(model_name, device=device)
model.set_pooling_include_prompt(include_prompt=False)
print(f"Model running on: {device}")

# ─────────────────────────────────────────────────────────────
# TEST MODE
# ─────────────────────────────────────────────────────────────
TEST_MODE    = False
TEST_SAMPLES = 100

# ─────────────────────────────────────────────────────────────
# 2. Dataset
# ─────────────────────────────────────────────────────────────
logging.info("Loading datasets...")

def load_csv(path):
    ds = load_dataset("csv", data_files=path)["train"]
    ds = ds.filter(lambda x: all(x[col] is not None and str(x[col]).strip() != "" for col in ds.column_names))
    if TEST_MODE:
        ds = ds.select(range(min(TEST_SAMPLES, len(ds))))
    return ds

train_data = load_csv("/home/skiredj.abderrahman/khalil/sbert_training/fourth_training/data/pre_train_arcd_SQuAD_triplets.csv")
val_data   = load_csv("/home/skiredj.abderrahman/khalil/sbert_training/fourth_training/data/pre_test_arcd_SQuAD_triplets.csv")

train_dataset = DatasetDict({"arcd_squad": train_data})
eval_dataset  = DatasetDict({"arcd_squad": val_data})

logging.info(train_dataset)
logging.info(eval_dataset)

# ─────────────────────────────────────────────────────────────
# 3. Loss
# ─────────────────────────────────────────────────────────────
loss = {
    "arcd_squad": losses.MatryoshkaLoss(
        model,
        losses.MultipleNegativesRankingLoss(model),
        matryoshka_dims=matryoshka_dims,
    )
}

# ─────────────────────────────────────────────────────────────
# 4. Evaluator
# ─────────────────────────────────────────────────────────────
def make_triplet_evaluators(dataset, name_prefix, max_samples=3_000):
    sample = dataset.shuffle(seed=42).select(range(min(max_samples, len(dataset))))
    return [
        TripletEvaluator(
            anchors=sample["anchor"],
            positives=sample["positive"],
            negatives=sample["negative"],
            name=f"{name_prefix}-{dim}",
            truncate_dim=dim,
        )
        for dim in matryoshka_dims
    ]

dev_evaluator = SequentialEvaluator(
    make_triplet_evaluators(val_data, "val-arcd-squad"),
    main_score_function=lambda scores: scores[0],
)

logging.info("Pre-training evaluation:")
dev_evaluator(model)

# ─────────────────────────────────────────────────────────────
# 5. Training Arguments
# ─────────────────────────────────────────────────────────────
args = SentenceTransformerTrainingArguments(
    output_dir=output_dir,
    seed=42,

    num_train_epochs=3,
    per_device_train_batch_size=train_batch_size,
    per_device_eval_batch_size=train_batch_size,
    gradient_accumulation_steps=2,

    bf16=True,
    fp16=False,
    learning_rate=1e-5,
    lr_scheduler_type="linear",
    warmup_ratio=0.1,
    weight_decay=0.01,

    batch_sampler=BatchSamplers.NO_DUPLICATES,
    multi_dataset_batch_sampler=MultiDatasetBatchSamplers.PROPORTIONAL,
    dataloader_num_workers=8,

    eval_strategy="steps",
    eval_steps=1000,
    save_strategy="steps",
    save_steps=1000,
    save_total_limit=2,

    report_to="tensorboard",
    logging_steps=100,
    logging_dir=f"{output_dir}/runs",
)

# ─────────────────────────────────────────────────────────────
# 6. Trainer
# ─────────────────────────────────────────────────────────────
trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=loss,
    evaluator=dev_evaluator,
)

trainer.train()
print("Finished")

# ─────────────────────────────────────────────────────────────
# 7. Save
# ─────────────────────────────────────────────────────────────
final_output_dir = f"{output_dir}/final"
model.save(final_output_dir)
print(f"Model saved to {final_output_dir}")