defmodule Bumblebee.Vision.ResNet do
  alias Bumblebee.Shared

  options =
    [
      num_channels: [
        default: 3,
        doc: "the number of channels in the input"
      ],
      embedding_size: [
        default: 64,
        doc: "the dimensionality of the embedding layer"
      ],
      hidden_sizes: [
        default: [256, 512, 1024, 2048],
        doc: "the dimensionality of hidden layers at each stage"
      ],
      depths: [
        default: [3, 4, 6, 3],
        doc: "the depth (number of residual blocks) at each stage"
      ],
      residual_block_type: [
        default: :bottleneck,
        doc: """
        the residual block to use, either `:basic` (used for smaller models, like ResNet-18 or ResNet-34)
        or `:bottleneck` (used for larger models like ResNet-50 and above)
        """
      ],
      activation: [
        default: :relu,
        doc: "the activation function"
      ],
      downsample_in_first_stage: [
        default: false,
        doc: "whether the first stage should downsample the inputs using a stride of 2"
      ]
    ] ++ Shared.common_options([:num_labels, :id_to_label])

  @moduledoc """
  ResNet model family.

  ## Architectures

    * `:base` - plain ResNet without any head on top

    * `:for_image_classification` - ResNet with a classification head.
      The head consists of a single dense layer on top of the pooled
      features and it returns logits corresponding to possible classes

  ## Inputs

    * `"pixel_values"` - {batch_size, height, width, num_channels}

      Featurized image pixel values (224x224).

  ## Global layer options

  #{Shared.global_layer_options_doc([:output_hidden_states])}

  ## Configuration

  #{Shared.options_doc(options)}
  """

  defstruct [architecture: :base] ++ Shared.option_defaults(options)

  @behaviour Bumblebee.ModelSpec
  @behaviour Bumblebee.Configurable

  import Bumblebee.Utils.Model, only: [join: 2]

  alias Bumblebee.Layers

  @impl true
  def architectures(), do: [:base, :for_image_classification]

  @impl true
  def config(spec, opts) do
    spec
    |> Shared.put_config_attrs(opts)
    |> Shared.validate_label_options()
  end

  @impl true
  def input_template(spec) do
    %{
      "pixel_values" => Nx.template({1, 224, 224, spec.num_channels}, :f32)
    }
  end

  @impl true
  def model(%__MODULE__{architecture: :base} = spec) do
    spec
    |> core()
    |> Layers.output()
  end

  def model(%__MODULE__{architecture: :for_image_classification} = spec) do
    outputs = core(spec)

    logits =
      Axon.dense(outputs.pooled_state, spec.num_labels, name: "image_classification_head.output")

    Layers.output(%{
      logits: logits,
      hidden_states: outputs.hidden_states
    })
  end

  defp core(spec, opts \\ []) do
    name = opts[:name]

    input = Axon.input("pixel_values", shape: {nil, 224, 224, spec.num_channels})

    encoder_outputs =
      input
      |> embedder(spec, name: join(name, "embedder"))
      |> encoder(spec, name: join(name, "encoder"))

    pooled_output =
      encoder_outputs.hidden_state
      |> Axon.adaptive_avg_pool(
        output_size: {1, 1},
        name: join(name, "pooler")
      )
      |> Axon.flatten()

    %{
      hidden_state: encoder_outputs.hidden_state,
      pooled_state: pooled_output,
      hidden_states: encoder_outputs.hidden_states
    }
  end

  defp embedder(pixel_values, spec, opts) do
    name = opts[:name]

    pixel_values
    |> conv_block(spec.embedding_size,
      kernel_size: 7,
      strides: 2,
      activation: spec.activation,
      name: join(name, "conv_block")
    )
    |> Axon.max_pool(
      kernel_size: 3,
      strides: 2,
      padding: [{1, 1}, {1, 1}],
      name: join(name, "pooler")
    )
  end

  defp encoder(hidden_state, spec, opts) do
    name = opts[:name]

    stages = spec.hidden_sizes |> Enum.zip(spec.depths) |> Enum.with_index()

    state = %{
      hidden_state: hidden_state,
      hidden_states: Axon.container({hidden_state}),
      in_channels: spec.embedding_size
    }

    for {{size, depth}, idx} <- stages, reduce: state do
      state ->
        strides = if idx == 0 and not spec.downsample_in_first_stage, do: 1, else: 2

        hidden_state =
          stage(state.hidden_state, state.in_channels, size, spec,
            depth: depth,
            strides: strides,
            name: join(name, "stages.#{idx}")
          )

        %{
          hidden_state: hidden_state,
          hidden_states: Layers.append(state.hidden_states, hidden_state),
          in_channels: size
        }
    end
  end

  defp stage(hidden_state, in_channels, out_channels, spec, opts) do
    opts = Keyword.validate!(opts, [:name, strides: 2, depth: 2])
    name = opts[:name]
    strides = opts[:strides]
    depth = opts[:depth]

    residual_block =
      case spec.residual_block_type do
        :basic -> &basic_residual_block/4
        :bottleneck -> &bottleneck_residual_block/4
      end

    # Downsampling is done in the first block with stride of 2
    hidden_state =
      residual_block.(hidden_state, in_channels, out_channels,
        strides: strides,
        activation: spec.activation,
        name: join(name, "blocks.0")
      )

    for idx <- 1..(depth - 1)//1, reduce: hidden_state do
      hidden_state ->
        residual_block.(hidden_state, out_channels, out_channels,
          activation: spec.activation,
          name: join(name, "blocks.#{idx}")
        )
    end
  end

  defp basic_residual_block(hidden_state, in_channels, out_channels, opts) do
    opts = Keyword.validate!(opts, [:name, strides: 1, activation: :relu])
    name = opts[:name]
    strides = opts[:strides]
    activation = opts[:activation]

    shortcut =
      shortcut(hidden_state, in_channels, out_channels,
        strides: strides,
        name: join(name, "shortcut")
      )

    hidden_state
    |> conv_block(out_channels, strides: strides, name: join(name, "conv_blocks.0"))
    |> conv_block(out_channels, activation: :linear, name: join(name, "conv_blocks.1"))
    |> Axon.add(shortcut)
    |> Axon.activation(activation, name: join(name, "activation"))
  end

  defp bottleneck_residual_block(hidden_state, in_channels, out_channels, opts) do
    opts = Keyword.validate!(opts, [:name, strides: 1, activation: :relu, reduction: 4])
    name = opts[:name]
    strides = opts[:strides]
    activation = opts[:activation]
    reduction = opts[:reduction]

    shortcut =
      shortcut(hidden_state, in_channels, out_channels,
        strides: strides,
        name: join(name, "shortcut")
      )

    reduced_channels = div(out_channels, reduction)

    hidden_state
    |> conv_block(reduced_channels, kernel_size: 1, name: join(name, "conv_blocks.0"))
    |> conv_block(reduced_channels, strides: strides, name: join(name, "conv_blocks.1"))
    |> conv_block(out_channels,
      kernel_size: 1,
      activation: :linear,
      name: join(name, "conv_blocks.2")
    )
    |> Axon.add(shortcut)
    |> Axon.activation(activation, name: join(name, "activation"))
  end

  defp shortcut(hidden_state, in_channels, out_channels, opts) do
    opts = Keyword.validate!(opts, [:name, strides: 2])
    name = opts[:name]
    strides = opts[:strides]

    # If the output shape doesn't match input shape, we need to project
    # the shortcut connection
    project_shortcut? = in_channels != out_channels or strides != 1

    if project_shortcut? do
      hidden_state
      |> Axon.conv(out_channels,
        kernel_size: 1,
        strides: strides,
        use_bias: false,
        kernel_initializer: conv_kernel_initializer(),
        name: join(name, "projection")
      )
      |> Axon.batch_norm(gamma_initializer: :ones, name: join(name, "norm"))
    else
      hidden_state
    end
  end

  defp conv_block(hidden_state, out_channels, opts) do
    opts = Keyword.validate!(opts, [:name, kernel_size: 3, strides: 1, activation: :relu])
    name = opts[:name]
    kernel_size = opts[:kernel_size]
    strides = opts[:strides]
    activation = opts[:activation]

    edge_padding = div(kernel_size, 2)
    padding_spec = [{edge_padding, edge_padding}, {edge_padding, edge_padding}]

    hidden_state
    |> Axon.conv(out_channels,
      kernel_size: kernel_size,
      strides: strides,
      padding: padding_spec,
      use_bias: false,
      kernel_initializer: conv_kernel_initializer(),
      name: join(name, "conv")
    )
    |> Axon.batch_norm(gamma_initializer: :ones, name: join(name, "norm"))
    |> Axon.activation(activation, name: join(name, "activation"))
  end

  defp conv_kernel_initializer() do
    Axon.Initializers.variance_scaling(scale: 2.0, mode: :fan_out, distribution: :normal)
  end

  defimpl Bumblebee.HuggingFace.Transformers.Config do
    def load(spec, data) do
      import Shared.Converters

      opts =
        convert!(data,
          num_channels: {"num_channels", number()},
          embedding_size: {"embedding_size", number()},
          hidden_sizes: {"hidden_sizes", list(number())},
          depths: {"depths", list(number())},
          residual_block_type: {"layer_type", atom()},
          activation: {"hidden_act", activation()},
          downsample_in_first_stage: {"downsample_in_first_stage", boolean()}
        ) ++ Shared.common_options_from_transformers(data, spec)

      @for.config(spec, opts)
    end
  end

  defimpl Bumblebee.HuggingFace.Transformers.Model do
    def params_mapping(_spec) do
      %{
        "embedder.conv_block.conv" => "resnet.embedder.embedder.convolution",
        "embedder.conv_block.norm" => "resnet.embedder.embedder.normalization",
        "encoder.stages.{n}.blocks.{m}.conv_blocks.{l}.conv" =>
          "resnet.encoder.stages.{n}.layers.{m}.layer.{l}.convolution",
        "encoder.stages.{n}.blocks.{m}.conv_blocks.{l}.norm" =>
          "resnet.encoder.stages.{n}.layers.{m}.layer.{l}.normalization",
        "encoder.stages.{n}.blocks.{m}.shortcut.projection" =>
          "resnet.encoder.stages.{n}.layers.{m}.shortcut.convolution",
        "encoder.stages.{n}.blocks.{m}.shortcut.norm" =>
          "resnet.encoder.stages.{n}.layers.{m}.shortcut.normalization",
        "image_classification_head.output" => "classifier.1"
      }
    end
  end
end
