import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer, DataCollatorForLanguageModeling
from transformers import Trainer, TrainingArguments
from datasets import load_dataset
import pandas as pd
import os

def check_device():
    """Check and return the available device"""
    if torch.cuda.is_available():
        print("Using GPU:", torch.cuda.get_device_name(0))
        return "cuda"
    else:
        print("No GPU available, using CPU")
        return "cpu"

def prepare_data(csv_file, text_column):
    """Prepare training data, filtering out empty rows."""
    df = pd.read_csv(csv_file)
    
    # Filter out rows where the specified column is NaN
    df = df.dropna(subset=[text_column])
    
    # Save texts to a file
    with open('train_data.txt', 'w', encoding='utf-8') as f:
        for text in df[text_column]:
            f.write(str(text) + '\n')
    
    return 'train_data.txt'

def tokenize_function(examples, tokenizer):
    """Tokenize the text using the GPT-2 tokenizer"""
    # Turn each sentence into codes (tokens) the robot can understand
    return tokenizer(examples['text'], padding="max_length", truncation=True, max_length=128)

def train_model(model, dataset, tokenizer):
    """Train the model with proper device configuration"""
    # Set up data collator
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False
    )
    
    # Create output directory if it doesn't exist
    os.makedirs("./gpt2-custom", exist_ok=True)
    
    # Set up training arguments with device-aware settings
    training_args = TrainingArguments(
        output_dir="./gpt2-custom",
        overwrite_output_dir=True,
        num_train_epochs=5,
        per_device_train_batch_size=3,  # Reduced batch size for CPU
        per_device_eval_batch_size=3,   # Reduced batch size for CPU
        save_steps=10_000,
        save_total_limit=2,
        logging_dir="./logs",
        logging_steps=500,
        use_cpu=(check_device() == "cpu"),  # Disable CUDA if no GPU
    )
    
    # Tokenize the dataset
    dataset = dataset.map(lambda examples: tokenize_function(examples, tokenizer), batched=True)
    dataset.set_format(type="torch", columns=["input_ids", "attention_mask"])

    # Initialize trainer with tokenized dataset
    trainer = Trainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=dataset['train'],
    )
    
    # Train the model
    trainer.train()
    
    # Save the model
    model.save_pretrained('./gpt2-custom')
    tokenizer.save_pretrained('./gpt2-custom')

def main():
    try:
        # Check device first
        device = check_device()
        
        # Example usage
        csv_file = './csvs/mission.csv'  # Replace with your CSV file path
        text_column = 'Mission'        # Replace with your text column name
        
        # Prepare data
        print("Preparing data...")
        train_file = prepare_data(csv_file, text_column)
        
        # Load model and tokenizer
        print("Loading model and tokenizer...")
        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        model = GPT2LMHeadModel.from_pretrained('gpt2')
        
        # Move model to appropriate device
        model = model.to(device)
        
        # Add padding token
        tokenizer.pad_token = tokenizer.eos_token
        
        # Load dataset with `datasets` library
        print("Loading dataset...")
        dataset = load_dataset('text', data_files={'train': train_file})

        # Train model
        print("Training model...")
        train_model(model, dataset, tokenizer)
        
        print("Training completed successfully!")
        
    except Exception as e:
        print(f"An error occurred: {str(e)}")
        print("Stack trace:")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()