import torch
import torch.nn as nn
from peft import get_peft_model, LoraConfig, TaskType

def finetune_with_lora(
    model,
    dataloader,
    epochs=1,
    lr=1e-4,
    device="cpu",
    r=8,
    alpha=16,
    dropout=0.05,
    target_modules=None,
    task_type=TaskType.CAUSAL_LM,
):
    model.to(device)
    model.train()

    config = LoraConfig(
        r=r,
        lora_alpha=alpha,
        target_modules=target_modules,
        lora_dropout=dropout,
        bias="none",
        task_type=task_type,
    )

    model = get_peft_model(model, config)
    model.print_trainable_parameters()

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

    for epoch in range(epochs):
        total_loss = 0.0
        for batch in dataloader:
            inputs = {k: v.to(device) for k, v in batch.items()}
            labels = inputs.get("input_ids")

            optimizer.zero_grad()
            outputs = model(**inputs)
            logits = outputs.logits
            loss = nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        print(f"[LoRA Fine-tune] Epoch {epoch+1}/{epochs} | Loss: {total_loss / len(dataloader):.4f}")

    return model