from src.data import PatchCamelyonDataset
from torch.utils.data import DataLoader

# Create dataset
train_dataset = PatchCamelyonDataset(
    root_dir="data/pcam",
    split="train"
)

# Create data loader
train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4
)

from src.models import SimpleClassifier
import torch

# Create model
model = SimpleClassifier(num_classes=2, dropout=0.5)

# Move to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

from torch import nn, optim

# Loss function
criterion = nn.CrossEntropyLoss()

# Optimizer
optimizer = optim.Adam(
    model.parameters(),
    lr=0.001,
    weight_decay=0.0001
)

# Learning rate scheduler
scheduler = optim.lr_scheduler.StepLR(
    optimizer,
    step_size=5,
    gamma=0.1
)

from src.training import train_epoch, evaluate
from src.utils import save_checkpoint

best_acc = 0.0

for epoch in range(10):
    # Train
    train_metrics = train_epoch(
        model, train_loader, criterion, optimizer, device, epoch
    )
    
    # Validate
    val_metrics = evaluate(
        model, val_loader, criterion, device
    )
    
    # Update learning rate
    scheduler.step()
    
    # Save best model
    if val_metrics['accuracy'] > best_acc:
        best_acc = val_metrics['accuracy']
        save_checkpoint(
            model, optimizer, epoch, val_metrics,
            "checkpoints/best_model.pth"
        )
    
    # Print progress
    print(f"Epoch {epoch+1}/10")
    print(f"  Train - Loss: {train_metrics['loss']:.4f}, Acc: {train_metrics['accuracy']:.4f}")
    print(f"  Val   - Loss: {val_metrics['loss']:.4f}, Acc: {val_metrics['accuracy']:.4f}")

print(f"Training complete! Best accuracy: {best_acc:.4f}")

from src.utils import load_checkpoint
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, roc_curve, auc
import seaborn as sns

# Load best model
checkpoint = load_checkpoint("checkpoints/best_model.pth", model)

# Evaluate
test_metrics = evaluate(model, test_loader, criterion, device)

# Confusion matrix
cm = confusion_matrix(test_metrics['labels'], test_metrics['predictions'])
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
plt.title('Confusion Matrix')
plt.ylabel('True Label')
plt.xlabel('Predicted Label')
plt.savefig('confusion_matrix.png')

# ROC curve
fpr, tpr, _ = roc_curve(test_metrics['labels'], test_metrics['probabilities'][:, 1])
roc_auc = auc(fpr, tpr)
plt.figure(figsize=(8, 6))
plt.plot(fpr, tpr, label=f'ROC curve (AUC = {roc_auc:.2f})')
plt.plot([0, 1], [0, 1], 'k--')
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('ROC Curve')
plt.legend()
plt.savefig('roc_curve.png')

print(f"Test Accuracy: {test_metrics['accuracy']:.4f}")
print(f"Test AUC: {roc_auc:.4f}")