"""Debug test - see if basic SFTTrainer creation works with Gemma4"""
import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model
from trl import SFTTrainer, SFTConfig

MODEL_ID = "google/gemma-4-E4B-it"
OUTPUT_DIR = "/tmp/gemma4-e4b-indic-lora"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
print(f"Tokenizer vocab size: {len(tokenizer)}")

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="eager",
    trust_remote_code=True,
)
model.config.use_cache = False
print(f"Model loaded. Params: {model.num_parameters():,}")

if torch.cuda.is_available():
    for i in range(torch.cuda.device_count()):
        mem = torch.cuda.memory_allocated(i) / 1024**3
        total = torch.cuda.get_device_properties(i).total_memory / 1024**3
        print(f"  GPU {i}: {mem:.2f}GB / {total:.2f}GB")

print("Unwrapping Gemma4ClippableLinear...")
count = 0
for name, module in model.named_modules():
    if module.__class__.__name__ == "Gemma4ClippableLinear":
        parts = name.split(".")
        parent_name = ".".join(parts[:-1]) if len(parts) > 1 else ""
        child_name = parts[-1]
        parent = model.get_submodule(parent_name) if parent_name else model
        setattr(parent, child_name, module.linear)
        count += 1
print(f"Unwrapped {count} layers")

print("Applying LoRA...")
peft_config = LoraConfig(
    r=32, lora_alpha=16, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()
print("LoRA applied")

print("Creating dummy dataset...")
from datasets import Dataset
dummy = Dataset.from_dict({"messages": [
    [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "4"}]
]})
print(f"Dummy dataset: {len(dummy)} samples")

print("\nCreating SFTTrainer...")
training_args = SFTConfig(
    output_dir=OUTPUT_DIR,
    max_length=128,
    packing=False,
    num_train_epochs=1,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=1,
    learning_rate=2e-4,
    bf16=True,
    logging_steps=1,
    disable_tqdm=False,
    max_steps=3,
    report_to="none",
    push_to_hub=False,
    dataloader_num_workers=0,
    remove_unused_columns=True,
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dummy,
    processing_class=tokenizer,
)
print("SFTTrainer created!")

print("\nStarting 3-step training...")
trainer.train()
print("Training complete!")