import gradio as gr
import torch
import torch.nn.functional as F
import nibabel as nib
import numpy as np

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

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

# Load the pre-trained model
model_path = "conv_autoencoder_model.pth"
model = ConvAutoencoder()
model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")), strict=False)
model.eval()

# Prediction function
def predict(modalities):
    try:
        slices = []
        for modality in modalities:
            nii_data = nib.load(modality).get_fdata()
            slices.append(nii_data)
        
        # Average modalities to create a single-channel image
        slice_avg = np.mean(slices, axis=0)
        
        # Select the slice with maximum non-zero pixels
        non_zero_counts = [np.count_nonzero(slice_avg.take(i, axis=2)) for i in range(slice_avg.shape[2])]
        max_index = np.argmax(non_zero_counts)
        best_slice = slice_avg.take(max_index, axis=2)
        
        # Convert best_slice to a tensor and resize to (256, 256) using interpolation
        tensor_slice = torch.tensor(best_slice).unsqueeze(0).unsqueeze(0).float()  # Shape: [1, 1, H, W]
        
        # Force resize to (256, 256)
        tensor_slice_resized = F.interpolate(tensor_slice, size=(256, 256), mode="bilinear", align_corners=False)
        print(f"Tensor shape after forced resize: {tensor_slice_resized.shape}")  # Debugging
        
        # Run model inference
        with torch.no_grad():
            reconstruction = model(tensor_slice_resized)
        reconstruction_error = torch.abs(reconstruction - tensor_slice_resized).squeeze().numpy()
        
        # Anomaly detection
        error_threshold = 0.1
        anomaly_detected = np.mean(reconstruction_error) > error_threshold
        anomaly_message = "Anomaly Detected" if anomaly_detected else "No Anomaly Detected"
        
        return reconstruction_error, anomaly_message

    except Exception as e:
        print(f"Processing error: {e}")
        return np.zeros((256, 256)), f"Processing error: {e}"

# Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Files(type="filepath", label="Upload NIfTI Modalities (e.g., T1, T2, FLAIR)"),
    outputs=[gr.Image(type="numpy", label="Reconstruction Error"), gr.Textbox(label="Anomaly Detection")],
    title="Brain MRI Anomaly Detection - ConvAutoencoder",
    description="Upload NIfTI files for brain anomaly detection using ConvAutoencoder.",
)

iface.launch()