from pathlib import Path

import torch
import torch.nn as nn
import torchio as tio

import matplotlib.pyplot as plt
import numpy as np

from auto_detect_breast_mri.config import resolve_path
from auto_detect_breast_mri.models.resnets import model_names
from auto_detect_breast_mri.data.breast_mri_dataset import BreastMRISubjects
# Paths come from the site config (see config.example.yaml).
path_base = resolve_path(None, "data_root", "root folder of the NIfTI data")
split_root = resolve_path(None, "split_root", "folder holding the split files")
model_root = resolve_path(None, "model_root", "folder holding the checkpoints", required=False) \
    or resolve_path(None, "output_root", "folder holding the checkpoints")

batch_size = 1
pre_image_shape = (32, 512, 512)
heatmap_size = (32, 256, 256)

# Get Model
model_name = 'resnet18_abrv'
protocol = ['Dyn_0', 'Sub_1']
model_path = str(Path(model_root) / "malignity_{}.pth".format(model_name))

model = model_names.get(model_name)

checkpoint = torch.load(model_path)
model.load_state_dict(checkpoint)
model.eval()

# Register forward hook
target_layer = model.layer4[-1]  # Choose the target layer
activation = {}


def get_activation(name):
    def hook(model, input, output):
        activation[name] = output.detach()

    return hook


target_layer.register_forward_hook(get_activation('target_layer'))

# forward pass
data_set = BreastMRISubjects(path_base, str(Path(split_root) / "fold0" / "stratified_evaluation_set-f0-0.csv"),
                           protocol=protocol)
data_loader = tio.SubjectsLoader(data_set, batch_size=batch_size, shuffle=False, num_workers=16)

for batch_id, batch in enumerate(data_loader):
    data = batch['image']['data']
    ground_truth = batch['label']
    path = batch['path']
    output = model(data)

    # Calculate gradients
    model.zero_grad()
    one_hot = torch.zeros_like(output)
    one_hot[0, output.argmax()] = 1.0
    output.backward(gradient=one_hot, retain_graph=True)

    # Weighted sum of feature maps
    grads = activation['target_layer'].grad.mean(dim=(2, 3), keepdim=True)
    heatmap = torch.sum(grads * activation['target_layer'], dim=1, keepdim=True)
    heatmap = nn.functional.relu(heatmap)

    # Resize the heatmap to match the original image size
    heatmap = nn.functional.interpolate(heatmap, size=heatmap_size, mode="bilinear", align_corners=False)

    heatmap = heatmap / heatmap.max()  # Normalize
    heatmap = heatmap.squeeze().numpy()  # transform to numpy array

    # Display the original image and Grad-CAM heatmap
    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    plt.imshow(np.transpose(data[0, 0, 16].squeeze().numpy(), (1, 2, 0)))
    plt.title('Original Image')

    plt.subplot(1, 2, 2)
    plt.imshow(heatmap[0, 0, 16], cmap='viridis')  # TODO maybe only 0, 16 ?
    plt.title('Grad-CAM Heatmap')

    plt.show()