from transformers import TrainingArguments, Trainer
from transformers import DataCollatorForSeq2Seq
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from datasets import load_dataset, load_from_disk
from textSummarizer.entity import ModelTrainerConfig
import torch 
import os

class ModelTrainer:
    def __init__(self, config: ModelTrainerConfig):
        self.config = config
        
    
    def train(self):
        device = "cuda" if torch.cuda.is_available() else "cpu"
        
        # Choose model checkpoint: prefer dev_model when dev_run is enabled
        model_checkpoint = str(self.config.model_ckpt)
        if getattr(self.config, 'dev_run', False) and getattr(self.config, 'dev_model', None):
            model_checkpoint = self.config.dev_model
            
        tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
        model_pegasus = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint).to(device)
        seq2seq_data_collator = DataCollatorForSeq2Seq(tokenizer, model=model_pegasus)
        
        dataset_samsum_pt = load_from_disk(str(self.config.data_path))
        
        # If dev_subset is set, use only a small subset of data for quick testing
        dev_subset = getattr(self.config, 'dev_subset', 0)
        if dev_subset > 0:
            if "train" in dataset_samsum_pt:
                dataset_samsum_pt["train"] = dataset_samsum_pt["train"].select(range(min(dev_subset, len(dataset_samsum_pt["train"]))))
            if "validation" in dataset_samsum_pt:
                validation_size = min(dev_subset // 4, len(dataset_samsum_pt["validation"]))
                dataset_samsum_pt["validation"] = dataset_samsum_pt["validation"].select(range(validation_size))

        # ensure datasets are returned as torch tensors for efficient dataloading
        if "train" in dataset_samsum_pt:
            try:
                dataset_samsum_pt["train"] = dataset_samsum_pt["train"].with_format("torch")
            except Exception:
                pass
        if "validation" in dataset_samsum_pt:
            try:
                dataset_samsum_pt["validation"] = dataset_samsum_pt["validation"].with_format("torch")
            except Exception:
                pass

        trainer_args = TrainingArguments(
            output_dir=str(self.config.root_dir),
            num_train_epochs=int(self.config.num_train_epochs),
            warmup_steps=int(self.config.warmup_steps),
            per_device_train_batch_size=int(self.config.per_device_train_batch_size),
            per_device_eval_batch_size=int(self.config.per_device_train_batch_size),
            weight_decay=float(self.config.weight_decay),
            logging_steps=int(self.config.logging_steps),
            eval_strategy=str(self.config.eval_strategy),
            eval_steps=int(self.config.eval_steps),
            save_steps=int(float(getattr(self.config, "save_steps", 1e6))),
            gradient_accumulation_steps=int(self.config.gradient_accumulation_steps),
            fp16=torch.cuda.is_available(),
            dataloader_num_workers=int(getattr(self.config, "dataloader_num_workers", 0)),
            save_total_limit=int(getattr(self.config, "save_total_limit", 1)),
            load_best_model_at_end=bool(getattr(self.config, "load_best_model_at_end", False)),
            # Adafactor uses ~4x less RAM than Adam — essential for Pegasus on CPU
            optim="adafactor",
        )

        trainer = Trainer(
            model=model_pegasus,
            args=trainer_args,
            data_collator=seq2seq_data_collator,
            train_dataset=dataset_samsum_pt["train"],
            eval_dataset=dataset_samsum_pt["validation"]
        )
        
        trainer.train()
        
        model_pegasus.save_pretrained(os.path.join(self.config.root_dir, "pegasus-samsum-model"))
        
        tokenizer.save_pretrained(os.path.join(self.config.root_dir, "tokenizer"))