import os
import torch
import sys
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling
)
from peft import LoraConfig, get_peft_model, TaskType

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import config

def main():
    print(f"Loading Base Model: {config.MODEL_ID}...")
    
    tokenizer = AutoTokenizer.from_pretrained(config.MODEL_ID, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        config.MODEL_ID,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True
    )
    
    # Configure LoRA specifically for Falcon Mamba target modules
    peft_config = LoraConfig(
        r=16,
        lora_alpha=32,
        target_modules=["in_proj", "x_proj", "dt_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type=TaskType.CAUSAL_LM
    )
    
    model = get_peft_model(model, peft_config)
    model.print_trainable_parameters()

    print(f"Loading datasets from {config.TRAIN_FILE} and {config.VAL_FILE}...")
    raw_dataset = load_dataset(
        "json",
        data_files={
            "train": config.TRAIN_FILE,
            "validation": config.VAL_FILE
        }
    )

    def tokenize_fn(examples):
        return tokenizer(
            examples["text"],
            truncation=True,
            max_length=1024,
            padding=False
        )

    print("Tokenizing train and validation splits...")
    tokenized_dataset = raw_dataset.map(
        tokenize_fn,
        batched=True,
        remove_columns=["text"]
    )

    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False
    )

    training_args = TrainingArguments(
        output_dir=config.OUTPUT_ADAPTER_DIR,
        per_device_train_batch_size=2,
        per_device_eval_batch_size=2,
        gradient_accumulation_steps=8,
        gradient_checkpointing=True,
        learning_rate=2e-4,
        weight_decay=0.01,
        num_train_epochs=3,
        logging_steps=10,
        eval_strategy="steps",
        eval_steps=100,
        save_strategy="steps",
        save_steps=200,
        save_total_limit=2,
        bf16=True,
        max_grad_norm=1.0,
        warmup_steps=50,
        lr_scheduler_type="cosine",
        report_to="none"
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset["train"],
        eval_dataset=tokenized_dataset["validation"],
        processing_class=tokenizer,
        data_collator=data_collator
    )

    print("Starting Mamba Reasoning fine-tuning on 2x NVIDIA H200 NVL GPUs...")
    trainer.train()
    
    print(f"Saving final fine-tuned adapter to {config.OUTPUT_ADAPTER_DIR}...")
    trainer.model.save_pretrained(config.OUTPUT_ADAPTER_DIR)
    tokenizer.save_pretrained(config.OUTPUT_ADAPTER_DIR)
    print("Fine-tuning completed successfully!")

if __name__ == "__main__":
    main()