import os
from typing import Optional, Tuple

import cv2
import numpy as np
import torch
import torch.nn as nn

from src.densenet import DenseNet121


def _preprocess_image(data_tensor: torch.Tensor) -> Tuple[np.ndarray, np.ndarray]:
    """
    Convert a (C, H, W) tensor into:
    - normalized float RGB image in [0, 1]
    - uint8 RGB image in [0, 255]
    """
    image = data_tensor.cpu().numpy()
    image = np.transpose(image, (1, 2, 0))
    image = (image - image.min()) / (image.max() - image.min())
    image_uint8 = (image * 255).astype(np.uint8)
    return image, image_uint8


def _process_smoothgrad_map(
    smoothgrad_map: np.ndarray,
    img_shape: Tuple[int, int],
    percentile: float,
    colormap: str,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Process raw SmoothGrad map into visualization-ready formats.
    """
    smoothgrad_map = np.transpose(smoothgrad_map, (1, 2, 0))
    smoothgrad_map = np.mean(smoothgrad_map, axis=2)
    smoothgrad_map = cv2.resize(smoothgrad_map, (img_shape[1], img_shape[0]))

    p_clip = np.percentile(smoothgrad_map, percentile)
    if p_clip > 0:
        smoothgrad_map = np.clip(smoothgrad_map, 0, p_clip)
        smoothgrad_map = (smoothgrad_map - smoothgrad_map.min()) / (p_clip - smoothgrad_map.min() + 1e-8)

    smoothgrad_heatmap = (smoothgrad_map * 255).astype(np.uint8)
    heatmap = cv2.applyColorMap(
        smoothgrad_heatmap,
        cv2.COLORMAP_HOT if colormap == "hot" else cv2.COLORMAP_VIRIDIS,
    )
    return smoothgrad_map, heatmap


def _get_target_layer(model: nn.Module, layer_name: Optional[str] = None):
    """
    Return DenseNet feature layer used for GradCAM.

    Default layer is `denseblock3` to reduce center-biased CAMs compared with
    the original `norm5_last`.

    You can override with:
    - function arg: `layer_name`
    - env var: `XDL_TARGET_LAYER`

    Supported layer names:
    `denseblock3`, `transition2`, `transition1`, `denseblock4`, `transition3`, `norm5_last`.
    """
    if not isinstance(model, DenseNet121):
        raise TypeError(
            f"Unsupported model type for this demo: {type(model).__name__}. Expected DenseNet121."
        )

    requested = (layer_name or os.getenv("XDL_TARGET_LAYER") or "denseblock3").strip().lower()
    aliases = {
        "default": "denseblock3",
        "last": "norm5_last",
        "norm5": "norm5_last",
    }
    selected = aliases.get(requested, requested)

    features = model.densenet_model.features
    layer_map = {
        "denseblock3": features.denseblock3,
        "transition2": features.transition2,
        "transition1": features.transition1,
        "denseblock4": features.denseblock4,
        "transition3": features.transition3,
        "norm5_last": features[-1],
    }

    if selected not in layer_map:
        supported = ", ".join(layer_map.keys())
        raise ValueError(
            f"Unsupported XDL target layer '{selected}'. Supported layers: {supported}"
        )

    return layer_map[selected]


def smoothgrad(
    model: nn.Module,
    input_tensor: torch.Tensor,
    target_class: int,
    n_samples: int = 100,
    noise_level: float = 0.05,
    batch_size: Optional[int] = None,
    use_amp: bool = False,
) -> np.ndarray:
    """
    Compute SmoothGrad saliency map for one input tensor of shape (1, C, H, W).

    Notes:
    - This implementation batches noisy samples to reduce per-step overhead.
    - If `input_tensor` and `model` are on CUDA, computation runs on GPU.
    """
    if n_samples <= 0:
        raise ValueError(f"n_samples must be > 0, got {n_samples}")

    if input_tensor.ndim != 4 or input_tensor.shape[0] != 1:
        raise ValueError(
            f"input_tensor must have shape (1, C, H, W), got {tuple(input_tensor.shape)}"
        )

    if batch_size is None:
        batch_size = min(16, n_samples)
    if batch_size <= 0:
        raise ValueError(f"batch_size must be > 0, got {batch_size}")

    model_param = next(model.parameters(), None)
    if model_param is not None and model_param.device != input_tensor.device:
        raise ValueError(
            f"Model device ({model_param.device}) and input device ({input_tensor.device}) must match."
        )

    model.eval()
    accumulated_gradients = torch.zeros_like(input_tensor)

    input_range = torch.max(input_tensor) - torch.min(input_tensor)
    scaled_noise = noise_level * input_range

    use_cuda_amp = use_amp and input_tensor.device.type == "cuda"
    for start in range(0, n_samples, batch_size):
        current_batch = min(batch_size, n_samples - start)
        expanded_input = input_tensor.expand(current_batch, -1, -1, -1)
        noise = torch.randn_like(expanded_input) * scaled_noise
        noisy_input = (expanded_input + noise).clone().detach().requires_grad_(True)

        if use_cuda_amp:
            with torch.autocast(device_type="cuda", dtype=torch.float16):
                output = model(noisy_input)
        else:
            output = model(noisy_input)
        if output.ndim == 1:
            output = output.unsqueeze(0)

        model.zero_grad(set_to_none=True)
        output[:, target_class].sum().backward()

        if noisy_input.grad is not None:
            accumulated_gradients += noisy_input.grad.data.sum(dim=0, keepdim=True)

    smooth = accumulated_gradients / n_samples
    smooth = torch.abs(smooth)
    smooth_min = torch.min(smooth)
    smooth_max = torch.max(smooth)
    smooth = (smooth - smooth_min) / (smooth_max - smooth_min + 1e-8)

    return smooth[0].cpu().numpy()