import os
from sklearn.model_selection import train_test_split
import numpy as np
import matplotlib.pyplot as plt


import logging
import sys
import traceback
from datetime import datetime

from datasets import load_dataset

from sentence_transformers import SentenceTransformer, losses
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator,TripletEvaluator,SequentialEvaluator
from sentence_transformers.similarity_functions import SimilarityFunction
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments
import torch
from sentence_transformers.training_args import BatchSamplers
import pandas as pd
from arabert.preprocess import ArabertPreprocessor
from pathlib import Path

# Configure logging to write to logs.txt
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

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

train_batch_size = 64

model_name = "bert-base-arabertv02"
model_nickname = "arabert"
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"

model = SentenceTransformer(model_name, device=device)
print(f"Model is running on: {device}")

logging.info("Reading the training and eval dataset")
train_dataset = load_dataset("csv", data_files="train.csv")
eval_dataset = load_dataset("csv", data_files="val.csv")
test_dataset = load_dataset("csv", data_files="test.csv")
logging.info(train_dataset)
logging.info(eval_dataset)
logging.info(test_dataset)
# Add this line before the 'evaluators = []' loop
eval_subset = eval_dataset["train"].shuffle(seed=42).select(range(250000))
# Training loss
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 for validation
evaluators = []
for dim in matryoshka_dims:
    evaluators.append(
        TripletEvaluator(
            anchors=eval_subset["anchor"],
            positives=eval_subset["positive"],
            negatives=eval_subset["negative"],
            name=f"dev-{dim}",
            truncate_dim=dim,
        )
    )
dev_evaluator = SequentialEvaluator(evaluators, main_score_function=lambda scores: scores[0])
dev_evaluator(model)


args = SentenceTransformerTrainingArguments(
    # --- Output & Identity ---
    output_dir=output_dir,  # Where the model and checkpoints are saved
    seed=42,                            # Ensures results are reproducible (shuffling/init)

    # --- Epochs & Batching (L40S Optimized) ---
    num_train_epochs=2,                 # Increased to 2 for better convergence
    per_device_train_batch_size=train_batch_size,    # High batch size to saturate the L40S 48GB VRAM
    per_device_eval_batch_size=train_batch_size,     # Matching eval batch size for speed
    gradient_accumulation_steps=2,      # Effective batch size = 128 (64 * 2)
    
    # --- Optimization & Precision ---
    bf16=True,                          # Set to True for L40S; faster and more stable than FP16
    fp16=False,                         # Disabled in favor of BF16
    learning_rate=2e-5,                 # Standard "safe" learning rate for Transformers
    lr_scheduler_type="linear",         # Gently reduces learning rate to 0 over training
    warmup_ratio=0.1,                   # Ramps up LR for the first 10% of steps to prevent spikes
    weight_decay=0.01,                  # Regularization to prevent overfitting

    # --- Data Handling ---
    batch_sampler=BatchSamplers.NO_DUPLICATES, # Essential for MNR/Matryoshka loss to avoid bad negatives
    dataloader_num_workers=8,           # Use 8 CPU cores to keep the GPU fed with data

    # --- Evaluation & Saving (Safety) ---
    eval_strategy="steps",              # Evaluate every X steps
    eval_steps=6000,                    # Increased from 10; L40S processes data very fast
    save_strategy="steps",              # Save checkpoints every X steps
    save_steps=6000,                    # Usually matches eval_steps
    save_total_limit=2,                 # Keep only the top 3 checkpoints to save disk space

    # --- Tracking & Debugging ---
    report_to="tensorboard",            # Send live metrics to TensorBoard
    logging_steps=200,                  # Print/Log stats every 200 steps (prevents messy logs)
    logging_dir="arabvert02-matryoshka/runs", # Specific folder for TensorBoard event files
)


trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=train_loss,
    evaluator=dev_evaluator,
)
trainer.train()

# Save final model
final_output_dir = f"{output_dir}/final"
model.save(final_output_dir)
print("model saved successfully")

# Test evaluation
evaluators = []
for dim in matryoshka_dims:
    evaluators.append(
        TripletEvaluator(
            anchors=test_dataset["train"]["anchor"],
            positives=test_dataset["train"]["positive"],
            negatives=test_dataset["train"]["negative"],
            name=f"test-{dim}",
            truncate_dim=dim,
        )
    )
test_evaluator = SequentialEvaluator(evaluators)
test_evaluator(model)