import os
import numpy as np
import cv2
from sklearn.model_selection import train_test_split
from transformers import ViTForImageClassification, ViTFeatureExtractor
from transformers import Trainer, TrainingArguments
from datasets import load_dataset, Dataset

# Function to load images from specified folders
def load_images_from_folders(folders, label):
    images = []
    labels = []
    for folder in folders:
        for filename in os.listdir(folder):
            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):  # Check for valid image extensions
                img = cv2.imread(os.path.join(folder, filename), cv2.IMREAD_GRAYSCALE)  # Read image as grayscale
                if img is not None:
                    img = cv2.resize(img, (224, 224))  # Resize to 224x224 pixels
                    img = img.astype(np.float32)  # Ensure the image is in float32 format
                    img /= 255.0  # Normalize to [0, 1]
                    images.append(img)
                    labels.append(label)
                else:
                    print(f"Failed to load image: {filename}")
    return images, labels

# Load normal and pneumonia images
normal_folders = [
    os.path.join('chest-xray-pneumonia', 'chest_xray', 'test', 'NORMAL'),
    os.path.join('chest-xray-pneumonia', 'chest_xray', 'train', 'NORMAL'),
    os.path.join('chest-xray-pneumonia', 'chest_xray', 'val', 'NORMAL'),
]
pneumonia_folders = [
    os.path.join('chest-xray-pneumonia', 'chest_xray', 'test', 'PNEUMONIA'),
    os.path.join('chest-xray-pneumonia', 'chest_xray', 'train', 'PNEUMONIA'),
    os.path.join('chest-xray-pneumonia', 'chest_xray', 'val', 'PNEUMONIA'),
]

normal_images, normal_labels = load_images_from_folders(normal_folders, 0)
pneumonia_images, pneumonia_labels = load_images_from_folders(pneumonia_folders, 1)

# Combine images and labels
images = normal_images + pneumonia_images
labels = normal_labels + pneumonia_labels

# Split the dataset into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(images, labels, test_size=0.2, random_state=42)

# Convert the dataset to a Hugging Face Dataset
train_dataset = Dataset.from_dict({"image": X_train, "label": y_train})
test_dataset = Dataset.from_dict({"image": X_test, "label": y_test})

# Load feature extractor and model
feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k', num_labels=2)

# Preprocess the dataset
def preprocess_function(examples):
    return feature_extractor(images=examples['image'], return_tensors="pt")

train_dataset = train_dataset.map(preprocess_function, batched=True)
test_dataset = test_dataset.map(preprocess_function, batched=True)

# Training arguments
training_args = TrainingArguments(
    output_dir='./results',
    evaluation_strategy='epoch',
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=10,
    weight_decay=0.01,
)

# Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
)

# Train the model
trainer.train()

# Save the model
model.save_pretrained('./pneumonia_model_final')
print("Model saved as './pneumonia_model_final'")