from transformers import TrainingArguments, Trainer, DetrForObjectDetection, DetrImageProcessor
import torch
import torchvision
from datasets import load_dataset, DatasetDict, Image

import os

# Set TRANSFORMERS_CACHE to a writable directory
os.environ["TRANSFORMERS_CACHE"] = "/app/cache"

# Ensure the directory exists
os.makedirs("/app/cache", exist_ok=True)

# Check if GPU is available
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
if torch.cuda.is_available():
    print(f"CUDA Device: {torch.cuda.get_device_name(0)}")

# Configuration
MODEL_NAME = "facebook/detr-resnet-50"  # DETR (DEtection TRansformer) model
BATCH_SIZE = 4
EPOCHS = 10
LEARNING_RATE = 5e-5
OUTPUT_DIR = "./results"

# Load Pre-trained Model and Processor
processor = DetrImageProcessor.from_pretrained(MODEL_NAME)
model = DetrForObjectDetection.from_pretrained(MODEL_NAME, num_labels=91).to(device)  # 91 is the number of COCO classes

# Load COCO Dataset
# Replace with paths to your custom dataset
coco_dataset = DatasetDict({
    "train": load_dataset("coco", data_files={"train": "dataset.coco.json"}, split="train"),
    "validation": load_dataset("coco", data_files={"validation": "dataset.coco.json"}, split="validation"),
})

# Preprocess Function for Images
def preprocess_coco(example):
    image = Image.open(example["file_name"]).convert("RGB")  # Ensure RGB format
    target = {
        "boxes": torch.tensor(example["bbox"]),
        "labels": torch.tensor(example["category_id"]),
    }
    encoding = processor(images=image, annotations=target, return_tensors="pt")
    return encoding

# Apply Preprocessing
coco_dataset = coco_dataset.map(preprocess_coco, batched=True)

# Data Collator
def collate_fn(batch):
    images = [item["pixel_values"].squeeze(0) for item in batch]
    annotations = [item["labels"] for item in batch]
    return {"pixel_values": torch.stack(images), "labels": annotations}

# Training Arguments
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    evaluation_strategy="epoch",
    learning_rate=LEARNING_RATE,
    per_device_train_batch_size=BATCH_SIZE,
    num_train_epochs=EPOCHS,
    save_strategy="epoch",
    logging_dir="./logs",
    logging_steps=10,
    load_best_model_at_end=True,
    push_to_hub=False,
)

# Trainer for Object Detection
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=coco_dataset["train"],
    eval_dataset=coco_dataset["validation"],
    tokenizer=processor,  # Not used for object detection, but required by Trainer
    data_collator=collate_fn,
)

# Train the Model
print("Starting Training...")
trainer.train()
print("Training Complete!")

# Save Model
print("Saving Model...")
trainer.save_model(OUTPUT_DIR)
print(f"Model saved to {OUTPUT_DIR}")