#' DeepLabV3 Semantic Segmentation Models
#'
#' @description
#' Semantic segmentation models implementing the DeepLabV3 architecture from
#' [Rethinking Atrous Convolution for Semantic Image Segmentation](https://arxiv.org/abs/1706.05587).
#' These models use Atrous Spatial Pyramid Pooling (ASPP) to capture multi-scale
#' context, and are available with ResNet-50 and ResNet-101 backbones.
#'
#' ## Available Models
#' - `model_deeplabv3_resnet50()`
#' - `model_deeplabv3_resnet101()`
#'
#' ## Model Variants and Performance (COCO val2017, VOC labels)
#'
#' All models are trained on a 20-class subset of COCO that corresponds to
#' Pascal VOC categories, plus background (21 classes total).
#'
#' ```
#' | Model                     | mIoU  | Pixel Acc | Params | GFLOPS | File Size | Weights Used              |
#' |---------------------------|-------|-----------|--------|--------|-----------|---------------------------|
#' | model_deeplabv3_resnet50  | 66.4% | 92.4%     | 42.0M  | 178.72 | 161 MB    | COCO_WITH_VOC_LABELS_V1   |
#' | model_deeplabv3_resnet101 | 67.4% | 92.4%     | 61.0M  | 258.74 | 233 MB    | COCO_WITH_VOC_LABELS_V1   |
#' ```
#'
#' ## Weights Selection
#' - All models use `COCO_WITH_VOC_LABELS_V1` weights, trained on COCO with the
#'   20 Pascal VOC categories (+ background = 21 classes).
#' - Backbone weights default to `IMAGENET1K_V1` (supervised ImageNet-1k) when
#'   `pretrained = FALSE` and `pretrained_backbone = TRUE`.
#' - When `pretrained = TRUE`, backbone weights are overridden by the full
#'   segmentation model weights and `pretrained_backbone` is ignored.
#' - The auxiliary classifier branch (`aux_loss`) is automatically enabled when
#'   loading pretrained weights; set explicitly when training from scratch.
#'
#' ## Input Format
#' Models expect input tensors of shape `(batch_size, 3, H, W)`, normalized
#' with ImageNet mean `c(0.485, 0.456, 0.406)` and std `c(0.229, 0.224, 0.225)`.
#' Training resolution is 520x520.
#'
#' ## Output Format
#' Returns a named list with:
#' - `$out` — main segmentation logits, shape `(batch, num_classes, H, W)`
#' - `$aux` — auxiliary logits from an intermediate backbone layer (only when `aux_loss = TRUE`)
#'
#' @inheritParams model_resnet18
#' @param num_classes Integer. Number of output segmentation classes including
#'   background. Default: `21` (Pascal VOC). Set to `NULL` to infer from
#'   pretrained weights.
#' @param aux_loss Logical or `NULL`. If `TRUE`, adds an auxiliary FCN classifier
#'   head at an intermediate backbone layer, used as a secondary loss during
#'   training. If `NULL` (default), inferred from pretrained weights.
#' @param pretrained_backbone Logical. If `TRUE` and `pretrained = FALSE`, loads
#'   `IMAGENET1K_V1` weights for the backbone only. Ignored when `pretrained = TRUE`.
#'   Default: `TRUE`.
#'
#' @examples
#' \dontrun{
#' library(magrittr)
#' norm_mean <- c(0.485, 0.456, 0.406)
#' norm_std  <- c(0.229, 0.224, 0.225)
#'
#' url <- paste0("https://upload.wikimedia.org/wikipedia/commons/thumb/",
#'            "e/ea/Morsan_Normande_vache.jpg/120px-Morsan_Normande_vache.jpg")
#' img <- base_loader(url)
#'
#' input <- img %>%
#'   transform_to_tensor() %>%
#'   transform_resize(c(520, 520)) %>%
#'   transform_normalize(norm_mean, norm_std)
#' batch <- input$unsqueeze(1)    # Add batch dimension: (1, 3, H, W)
#'
#' # --- ResNet-50 backbone ---
#' model <- model_deeplabv3_resnet50(pretrained = TRUE)
#' model$eval()
#' output <- model(batch)
#'
#' segmented <- draw_segmentation_masks(input, output$out$squeeze(1))
#' tensor_image_browse(segmented)
#'
#' # Show most frequent class
#' mask_id <- output$out$argmax(dim = 2)  # (1, H, W)
#' class_contingency_with_background <- mask_id$view(-1)$bincount()
#' class_contingency_with_background[1] <- 0L # we clean the counter for background class id 1
#' top_class_index <- class_contingency_with_background$argmax()$item()
#' cli::cli_inform("Majority class {.pkg ResNet-50}: {.emph {pascal_voc_classes(top_class_index)}}")
#'
#' # --- ResNet-101 backbone ---
#' model <- model_deeplabv3_resnet101(pretrained = TRUE)
#' model$eval()
#' output <- model(batch)
#'
#' segmented <- draw_segmentation_masks(input, output$out$squeeze(1))
#' tensor_image_browse(segmented)
#'
#' # Show most frequent class
#' mask_id <- output$out$argmax(dim = 2)  # (1, H, W)
#' class_contingency_with_background <- mask_id$view(-1)$bincount()
#' class_contingency_with_background[1] <- 0L # we clean the counter for background class id 1
#' top_class_index <- class_contingency_with_background$argmax()$item()
#' cli::cli_inform("Majority class {.pkg ResNet-50}: {.emph {pascal_voc_classes(top_class_index)}}")
#' }
#'
#' @importFrom torch nn_module
#' @family semantic_segmentation_model
#' @rdname model_deeplabv3
#' @name model_deeplabv3
NULL

deeplabv3_model_urls <- list(
  deeplabv3_resnet50_coco = c(
    "https://torch-cdn.mlverse.org/models/vision/v2/models/deeplabv3_resnet50_coco.pth",
    "9bc0953c76291fff296211b0eba04164",
    "168 MB"
  ),
  deeplabv3_resnet101_coco = c(
    "https://torch-cdn.mlverse.org/models/vision/v2/models/deeplabv3_resnet101_coco.pth",
    "bb73d96704937621d782af73125e6282",
    "172 MB"
  )
)

deeplabv3_meta <- list(
  classes = pascal_voc_classes(),
  class_to_idx = setNames(seq_along(pascal_voc_classes()) - 1, pascal_voc_classes())
)


# ASPP module matching PyTorch
aspp_module <- torch::nn_module(
  "ASPP",
  initialize = function(in_channels, out_channels, atrous_rates) {
    self$convs <- torch::nn_module_list()

    self$convs$append(torch::nn_sequential(
      torch::nn_conv2d(in_channels, out_channels, kernel_size = 1, bias = FALSE),
      torch::nn_batch_norm2d(out_channels),
      torch::nn_relu()
    ))

    for (rate in atrous_rates) {
      self$convs$append(torch::nn_sequential(
        torch::nn_conv2d(in_channels, out_channels, kernel_size = 3, padding = rate, dilation = rate, bias = FALSE),
        torch::nn_batch_norm2d(out_channels),
        torch::nn_relu()
      ))
    }

    self$convs$append(torch::nn_sequential(
      torch::nn_adaptive_avg_pool2d(output_size = c(1, 1)),
      torch::nn_conv2d(in_channels, out_channels, kernel_size = 1, bias = FALSE),
      torch::nn_batch_norm2d(out_channels),
      torch::nn_relu()
    ))

    self$project <- torch::nn_sequential(
      torch::nn_conv2d((length(atrous_rates) + 2) * out_channels, out_channels, kernel_size = 1, bias = FALSE),
      torch::nn_batch_norm2d(out_channels),
      torch::nn_relu(),
      torch::nn_dropout(0.5)
    )
  },
  forward = function(x) {
    input_size <- x$shape[3:4]  # Get height and width dimensions
    res <- list()

    for (i in 1:(length(self$convs) - 1)) {
      res[[i]] <- self$convs[[i]](x)
    }

    global_feat <- self$convs[[length(self$convs)]](x)
    target_size <- as.integer(input_size)
    global_feat <- nnf_interpolate(global_feat, size = target_size, mode = "bilinear", align_corners = FALSE)
    res[[length(res) + 1]] <- global_feat

    x <- torch_cat(res, dim = 2)
    self$project(x)
  }
)

# Main classifier head
deeplab_head <- function(in_channels, num_classes) {
  torch::nn_sequential(
    aspp_module(in_channels, 256, atrous_rates = c(12, 24, 36)),
    torch::nn_conv2d(256, 256, kernel_size = 3, padding = 1, bias = FALSE),
    torch::nn_batch_norm2d(256),
    torch::nn_relu(),
    torch::nn_conv2d(256, num_classes, kernel_size = 1)
  )
}

# Auxiliary classifier head (simpler, operates on layer3 features)
aux_classifier_head <- function(in_channels, num_classes) {
  torch::nn_sequential(
    torch::nn_conv2d(in_channels, 256, kernel_size = 3, padding = 1, bias = FALSE),
    torch::nn_batch_norm2d(256),
    torch::nn_relu(),
    torch::nn_dropout(0.1),
    torch::nn_conv2d(256, num_classes, kernel_size = 1)
  )
}

# Updated DeepLabV3 wrapper with auxiliary classifier
DeepLabV3 <- torch::nn_module(
  "DeepLabV3",
  initialize = function(backbone, classifier, aux_classifier = NULL) {
    self$backbone <- backbone
    self$classifier <- classifier
    self$aux_classifier <- aux_classifier
  },
  forward = function(x) {
    input_shape <- x$shape[3:4]  # Get height and width dimensions

    x <- self$backbone$conv1(x)
    x <- self$backbone$bn1(x)
    x <- self$backbone$relu(x)
    x <- self$backbone$maxpool(x)
    x <- self$backbone$layer1(x)
    x <- self$backbone$layer2(x)
    x <- self$backbone$layer3(x)

    aux_x <- x

    x <- self$backbone$layer4(x)

    # Main classifier
    x <- self$classifier(x)
    target_size <- as.integer(input_shape)
    x <- nnf_interpolate(x, size = target_size, mode = "bilinear", align_corners = FALSE)

    result <- list(out = x)

    # Auxiliary classifier (if present)
    if (!is.null(self$aux_classifier)) {
      aux_out <- self$aux_classifier(aux_x)
      aux_out <- nnf_interpolate(aux_out, size = target_size, mode = "bilinear", align_corners = FALSE)
      result$aux <- aux_out
    }

    result
  }
)

# Model constructor
deeplabv3_resnet_factory <- function(arch, block, layers, pretrained, progress,
                                     num_classes, aux_loss = NULL,
                                     pretrained_backbone = FALSE, ...) {

  validate_num_classes(num_classes, pretrained)

  if (is.null(aux_loss))
    aux_loss <- FALSE

  if (pretrained && pretrained_backbone)
    cli_warn("`pretrained_backbone` ignored when `pretrained = TRUE`." )

  backbone_arch <- sub("deeplabv3_", "", arch)

  if (pretrained_backbone && !pretrained) {
    backbone <- .resnet(backbone_arch, block, layers, pretrained = TRUE,
                        progress = progress,
                        replace_stride_with_dilation = c(FALSE, TRUE, TRUE), ...)
  } else {
    backbone <- resnet(block, layers,
                       replace_stride_with_dilation = c(FALSE, TRUE, TRUE), ...)
  }
  backbone$fc <- torch::nn_identity()
  backbone$avgpool <- torch::nn_identity()

  classifier <- deeplab_head(2048, num_classes)

  aux_classifier <- NULL
  if (aux_loss) {
    aux_classifier <- aux_classifier_head(1024, num_classes)
  }

  model <- DeepLabV3(backbone, classifier, aux_classifier)
  model$classes <- deeplabv3_meta$classes
  model$class_to_idx <- deeplabv3_meta$class_to_idx

  if (pretrained) {
    info <- deeplabv3_model_urls[[paste0(arch, "_coco")]]
    cli_inform("Downloading {.cls {arch}} pretrained weights (~{.emph {info[3]}}) ...")
    state_dict_path <- download_and_cache(info[1], prefix = "deeplabv3")
    if (tools::md5sum(state_dict_path) != info[2])
      runtime_error("Corrupt file! Delete the file in {state_dict_path} and try again.")
    state_dict <- torch::load_state_dict(state_dict_path)
    model$load_state_dict(state_dict)
  }
  model
}

#' @describeIn model_deeplabv3 DeepLabV3 with ResNet-50 backbone
#' @export
model_deeplabv3_resnet50 <- function(pretrained = FALSE, progress = TRUE,
                                     num_classes = 21, aux_loss = NULL,
                                     pretrained_backbone = FALSE, ...) {
  deeplabv3_resnet_factory(
    "deeplabv3_resnet50", bottleneck, c(3, 4, 6, 3),
    pretrained, progress, num_classes, aux_loss, pretrained_backbone, ...
  )
}

#' @describeIn model_deeplabv3 DeepLabV3 with ResNet-101 backbone
#' @export
model_deeplabv3_resnet101 <- function(pretrained = FALSE, progress = TRUE,
                                      num_classes = 21, aux_loss = NULL,
                                      pretrained_backbone = FALSE, ...) {
  deeplabv3_resnet_factory(
    "deeplabv3_resnet101", bottleneck, c(3, 4, 23, 3),
    pretrained, progress, num_classes, aux_loss, pretrained_backbone, ...
  )
}