import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from PIL import Image

# Define the ConvAutoencoder model structure
class ConvAutoencoder(nn.Module):
    def __init__(self):
        super(ConvAutoencoder, self).__init__()
        # Encoder
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 32, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.ReLU()
        )
        # Decoder
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 32, 2, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 1, 2, stride=2, padding=1),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

# Load and resize the training slice
processed_data_path = 'BraTS20_Training_001_best_slice.npy'  # Replace with your actual path
best_slice = np.load(processed_data_path)

# Resize the best slice to (256, 256)
transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((256, 256)),  # Ensures fixed input size
    transforms.ToTensor()
])

# Apply transformation to ensure correct size
best_slice = transform(best_slice).unsqueeze(0)  # Shape: [1, 1, 256, 256]

# Initialize model, loss function, and optimizer
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ConvAutoencoder().to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# Training loop
num_epochs = 100
best_slice = best_slice.to(device)

for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()
    output = model(best_slice)
    loss = criterion(output, best_slice)
    loss.backward()
    optimizer.step()
    
    if (epoch + 1) % 10 == 0:
        print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}")

# Save the trained model weights
torch.save(model.state_dict(), "conv_autoencoder_model.pth")
print("Model training complete and saved as conv_autoencoder_model.pth")