#30. ToDo

import torch
import os
from transformers import AutoModelForImageClassification, AutoImageProcessor
from datasets import load_dataset
from torch.utils.data import DataLoader, Dataset
from torch.optim import Adam
import torch.nn as nn
from sklearn.preprocessing import LabelEncoder
from peft import get_peft_model, LoraConfig
import requests
import io
from torchvision import transforms
from tqdm import tqdm
from sklearn.model_selection import train_test_split

# ViT 모델 로드 및 전처리 함수
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224", use_fast=True)
model = AutoModelForImageClassification.from_pretrained(
    "google/vit-base-patch16-224",
    num_labels=8,
    ignore_mismatched_sizes=True
).to(device)

# 전처리 함수
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
    ])
    return transform(image)

class CustomDataset(Dataset):
    def __init__(self, dataset, transform=None):
        self.dataset = dataset
        self.transform = transform
        self.label_encoder = LabelEncoder()
        labels = [item['label'] for item in dataset]
        self.label_encoder.fit(labels)

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        item = self.dataset[idx]
        img = item['image']
        label = item['label']
        if self.transform:
            img = self.transform(img)
        label = self.label_encoder.transform([label])[0]
        return img, torch.tensor(label, dtype=torch.long)

# 모델 가중치 로드 및 데이터셋 확인
model_url = "https://huggingface.co/JANGJIWON/UGRP1.0/resolve/main/model_epoch_3_accuracy_23.17.pth"
response = requests.get(model_url)
model_weights = io.BytesIO(response.content)

# 첫 번째 LoRA 구성 및 적용
config1 = LoraConfig(
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
    target_modules=["query", "key", "value"]
)
model = get_peft_model(model, config1)

try:
    model.load_state_dict(torch.load(model_weights, map_location='cpu', weights_only=True), strict=False)
except RuntimeError as e:
    print(f"Error loading state_dict: {e}")

# 두 번째 LoRA 구성 및 적용
config2 = LoraConfig(
    r=4,
    lora_alpha=16,
    lora_dropout=0.05,
    target_modules=["query", "key", "value"]
)
model = get_peft_model(model, config2)

# 데이터셋 준비
dataset = load_dataset("JANGJIWON/test", split="train")

# Convert the dataset to a list of dictionaries for splitting
dataset_list = [dict(item) for item in dataset]

# Split the dataset into train and test sets (80% train, 20% test)
train_data, test_data = train_test_split(dataset_list, test_size=0.2, random_state=42)

# Create datasets and dataloaders
train_dataset = CustomDataset(train_data, transform=preprocess_image)
test_dataset = CustomDataset(test_data, transform=preprocess_image)

train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False, num_workers=2)

# 옵티마이저 및 손실 함수 설정
learning_rate = 1  # 조금 더 높은 학습률로 설정
optimizer = Adam(model.parameters(), lr=learning_rate)
criterion = nn.CrossEntropyLoss()

# 모델 훈련
num_epochs = 3
model.train()

for epoch in range(num_epochs):
    running_loss = 0.0
    correct = 0
    total = 0
    for images, labels in tqdm(train_loader):
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()

        outputs = model(images)
        loss = criterion(outputs.logits, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = torch.max(outputs.logits, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    epoch_accuracy = 100 * correct / total
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss:.4f}, Accuracy: {epoch_accuracy:.2f}%')

# 테스트 수행 및 정확도 계산
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:  # Use test_loader here
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.logits, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f'Final Test Accuracy after second LoRA tuning: {accuracy:.2f}%')