import torch
from torch.utils.data import DataLoader, random_split
from datasets import load_dataset
from src.dataset import HumanActionDataset
from utils.preprocessing import get_transforms
from models.resnet_model import ResNet18
import pandas as pd

def get_val_loader(batch_size=32):
    ds = load_dataset("Bingsu/Human_Action_Recognition")
    full_dataset = HumanActionDataset(ds['train'], transform=get_transforms())

    train_size = int(0.8 * len(full_dataset))
    val_size = len(full_dataset) - train_size
    _, val_dataset = random_split(full_dataset, [train_size, val_size])

    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
    return val_loader

def get_test_loader(batch_size=32):
    ds = load_dataset("Bingsu/Human_Action_Recognition")
    test_dataset = HumanActionDataset(ds['test'], transform=get_transforms())
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    return test_loader

def evaluate_model(model_path, batch_size=32):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    model = ResNet18(num_classes=15)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    model.eval()

    val_loader = get_val_loader(batch_size)

    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in val_loader:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            _, predicted = torch.max(outputs, 1)

            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = correct / total
    print(f"Validation Accuracy: {accuracy:.4f} ({correct}/{total})")

def predict_test(model_path, batch_size=32, output_csv="test_predictions.csv"):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    model = ResNet18(num_classes=15)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    model.eval()

    test_loader = get_test_loader(batch_size)
    class_names = [
        'calling', 'clapping', 'cycling', 'dancing', 'drinking', 'eating', 'fighting',
        'hugging', 'laughing', 'listening_to_music', 'running', 'sitting', 'sleeping',
        'texting', 'using_laptop'
    ]

    predictions = []
    filenames = []

    with torch.no_grad():
        for batch in test_loader:
            images = batch[0].to(device)
            # Assuming your HumanActionDataset returns (image, label, filename) or just (image, label)
            # If filenames are needed and your dataset doesn't provide, you can omit or modify accordingly.
            
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)

            predictions.extend(predicted.cpu().tolist())
            # If you have filenames:
            # filenames.extend(batch[2])  # Uncomment if your dataset yields filenames

    # If you don't have filenames, just save predictions indexed by order
    df = pd.DataFrame({
        "id": list(range(len(predictions))),
        "label": [class_names[p] for p in predictions]
    })
    df.to_csv(output_csv, index=False)
    print(f"Saved test predictions to {output_csv}")

if __name__ == "__main__":
    model_path = "models/best_model.pth"
    evaluate_model(model_path=model_path, batch_size=32)
    predict_test(model_path=model_path, batch_size=32, output_csv="test_predictions.csv")