defmodule Bumblebee.Vision.DinoV2 do
  alias Bumblebee.Shared

  options =
    [
      image_size: [
        default: 518,
        doc: """
        the size of the input spatial dimensions. The model is trained for this size, however
        the model supports any other input size by interpolating position embeddings
        """
      ],
      num_channels: [
        default: 3,
        doc: "the number of channels in the input"
      ],
      patch_size: [
        default: 14,
        doc: "the size of the patch spatial dimensions"
      ],
      hidden_size: [
        default: 384,
        doc: "the dimensionality of hidden layers"
      ],
      num_blocks: [
        default: 12,
        doc: "the number of Transformer blocks in the encoder"
      ],
      num_attention_heads: [
        default: 12,
        doc: "the number of attention heads for each attention layer in the encoder"
      ],
      intermediate_size_ratio: [
        default: 4,
        doc: """
        the dimensionality of the intermediate layer in the transformer feed-forward network (FFN) in the encoder,
        expressed as a multiplier of `:hidden_size`
        """
      ],
      use_attention_bias: [
        default: true,
        doc: "whether to use bias in query, key, and value projections"
      ],
      activation: [
        default: :gelu,
        doc: "the activation function"
      ],
      ffn_swiglu_activation: [
        default: false,
        doc:
          "whether to use the gated SwiGLU activation function in the feed-forward network (FFN)"
      ],
      scale_initial_value: [
        default: 1.0,
        doc: "the initial value for scaling layers"
      ],
      dropout_rate: [
        default: 0.0,
        doc: "the dropout rate for encoder and decoder"
      ],
      attention_dropout_rate: [
        default: 0.0,
        doc: "the dropout rate for attention weights"
      ],
      layer_norm_epsilon: [
        default: 1.0e-6,
        doc: "the epsilon used by the layer normalization layers"
      ],
      initializer_scale: [
        default: 0.02,
        doc:
          "the standard deviation of the normal initializer used for initializing kernel parameters"
      ],
      backbone_output_indices: [
        default: nil,
        doc: """
        list of indices indicating which feature maps to include in the output. If not specified, only
        the last feature map is included
        """
      ],
      backbone_use_norm: [
        default: true,
        doc:
          "whether to add layer normalization layer to each of the feature maps returned by the backbone"
      ]
    ] ++ Shared.common_options([:num_labels, :id_to_label])

  @moduledoc """
  DINOv2 model family.

  ## Architectures

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

    * `:for_image_classification` - DINOv2 with head for image classification

    * `:backbone` - DINOv2 with feature maps output

  ## Inputs

    * `"pixel_values"` - `{batch_size, image_size, image_size, num_channels}`

      Featurized image pixel values.

    * `"patch_mask"` - `{batch_size, num_patches}`

      Mask to nullify selected embedded patches.

  ## Global layer options

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

  ## Configuration

  #{Shared.options_doc(options)}

  ## References

    * [DINOv2: Learning Robust Visual Features without Supervision](https://arxiv.org/abs/2304.07193)

  """

  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, :backbone, :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, spec.image_size, spec.image_size, spec.num_channels}, :f32)
    }
  end

  @impl true
  def model(%__MODULE__{architecture: :base} = spec) do
    inputs = inputs(spec)
    outputs = core(inputs, spec)

    hidden_state =
      Axon.layer_norm(outputs.hidden_state, epsilon: spec.layer_norm_epsilon, name: "norm")

    pooled_state = Layers.take_token(hidden_state, index: 0, axis: 1)

    Layers.output(%{
      hidden_state: hidden_state,
      pooled_state: pooled_state,
      hidden_states: outputs.hidden_states,
      attentions: outputs.attentions
    })
  end

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

    hidden_state =
      Axon.layer_norm(outputs.hidden_state, epsilon: spec.layer_norm_epsilon, name: "norm")

    class_token = Layers.take_token(hidden_state, index: 0, axis: 1)

    patch_embeddings_mean =
      Axon.nx(hidden_state, fn hidden_state ->
        patch_embeddings = hidden_state[[.., 1..-1//1, ..]]
        Nx.mean(patch_embeddings, axes: [1])
      end)

    logits =
      Axon.concatenate(class_token, patch_embeddings_mean)
      |> Axon.dense(spec.num_labels,
        kernel_initializer: kernel_initializer(spec),
        name: "image_classification_head.output"
      )

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

  def model(%__MODULE__{architecture: :backbone} = spec) do
    inputs = inputs(spec)
    outputs = core(inputs, spec)

    feature_maps = feature_maps(outputs.hidden_states, inputs["pixel_values"], spec)

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

  defp inputs(spec) do
    shape = {nil, nil, nil, spec.num_channels}

    Bumblebee.Utils.Model.inputs_to_map([
      Axon.input("pixel_values", shape: shape),
      Axon.input("patch_mask", shape: {nil, nil}, optional: true)
    ])
  end

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

    embeddings =
      embedder(inputs["pixel_values"], inputs["patch_mask"], spec, name: join(name, "embedder"))

    encoder_outputs = encoder(embeddings, spec, name: join(name, "encoder"))

    %{
      hidden_state: encoder_outputs.hidden_state,
      hidden_states: encoder_outputs.hidden_states,
      attentions: encoder_outputs.attentions
    }
  end

  defp feature_maps(hidden_states, pixel_values, spec, opts \\ []) do
    name = opts[:name]

    num_feature_maps = spec.num_blocks + 1
    output_indices = spec.backbone_output_indices || [num_feature_maps - 1]

    for index <- output_indices do
      hidden_state = Axon.nx(hidden_states, &elem(&1, index))

      hidden_state =
        if spec.backbone_use_norm do
          Axon.layer_norm(hidden_state,
            epsilon: spec.layer_norm_epsilon,
            name: join(name, "norm")
          )
        else
          hidden_state
        end

      Axon.layer(
        fn hidden_state, pixel_values, _opts ->
          {batch_size, height, width, _channels} = Nx.shape(pixel_values)

          hidden_state = hidden_state[[.., 1..-1//1, ..]]

          Nx.reshape(
            hidden_state,
            {batch_size, div(height, spec.patch_size), div(width, spec.patch_size), :auto}
          )
        end,
        [hidden_state, pixel_values]
      )
    end
    |> List.to_tuple()
    |> Axon.container()
  end

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

    patch_embeddings =
      pixel_values
      |> patch_embedding(spec, name: join(name, "patch_embedding"))
      |> Layers.apply_vision_patch_mask(patch_mask, name: join(name, "mask_tokens"))

    class_embedding =
      Layers.learned_embeddings(1, spec.hidden_size, name: join(name, "class_embedding"))

    input_embeddings = Layers.concatenate_embeddings([class_embedding, patch_embeddings])

    num_patches = div(spec.image_size, spec.patch_size) ** 2

    position_embeddings =
      Layers.learned_embeddings(num_patches + 1, spec.hidden_size,
        initializer: :zeros,
        name: join(name, "position_embedding")
      )
      |> interpolate_position_embeddings(pixel_values, spec)

    Axon.add(input_embeddings, position_embeddings)
    |> Axon.dropout(rate: spec.dropout_rate, name: join(name, "dropout"))
  end

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

    pixel_values
    |> Axon.conv(spec.hidden_size,
      kernel_size: spec.patch_size,
      strides: spec.patch_size,
      padding: :valid,
      kernel_initializer: kernel_initializer(spec),
      name: join(name, "projection")
    )
    |> Axon.reshape({:batch, :auto, spec.hidden_size}, name: join(name, "reshape"))
  end

  defp interpolate_position_embeddings(position_embeddings, pixel_values, spec) do
    Axon.layer(
      fn position_embeddings, pixel_values, _opts ->
        original_positions = div(spec.image_size, spec.patch_size)
        {batch_size, height, width, _channels} = Nx.shape(pixel_values)
        resized_height = div(height, spec.patch_size)
        resized_width = div(width, spec.patch_size)

        class_position_embedding = position_embeddings[[.., 0..0//1, ..]]
        input_position_embeddings = position_embeddings[[.., 1..-1//1, ..]]

        interpolated_position_embeddings =
          input_position_embeddings
          |> Nx.reshape({batch_size, original_positions, original_positions, spec.hidden_size})
          |> Axon.Layers.resize(
            size: {resized_height, resized_width},
            method: :bicubic,
            antialias: false
          )
          |> Nx.reshape({batch_size, :auto, spec.hidden_size})

        Nx.concatenate([class_position_embedding, interpolated_position_embeddings], axis: 1)
      end,
      [position_embeddings, pixel_values],
      op_name: :interpolate_position_embeddings
    )
  end

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

    ffn =
      if spec.ffn_swiglu_activation do
        intermediate_size =
          div(floor(floor(spec.hidden_size * spec.intermediate_size_ratio) * 2 / 3 + 7), 8) * 8

        &ffn_swiglu(&1, intermediate_size, spec.hidden_size, name: &2)
      else
        intermediate_size = floor(spec.hidden_size * spec.intermediate_size_ratio)

        [
          intermediate_size: intermediate_size,
          activation: spec.activation
        ]
      end

    Layers.Transformer.blocks(hidden_state,
      num_blocks: spec.num_blocks,
      num_attention_heads: spec.num_attention_heads,
      hidden_size: spec.hidden_size,
      kernel_initializer: kernel_initializer(spec),
      dropout_rate: spec.dropout_rate,
      attention_dropout_rate: spec.attention_dropout_rate,
      query_use_bias: spec.use_attention_bias,
      key_use_bias: spec.use_attention_bias,
      value_use_bias: spec.use_attention_bias,
      layer_norm: [
        epsilon: spec.layer_norm_epsilon
      ],
      ffn: ffn,
      block_type: &block_impl(&1, &2, &3, spec),
      name: join(name, "blocks")
    )
  end

  # A feed-forward network with SwiGLU nonlinearity as in https://arxiv.org/abs/2002.05202
  defp ffn_swiglu(x, intermediate_size, output_size, opts) do
    name = opts[:name]
    dropout = opts[:dropout] || 0.0

    {gate, x} =
      x
      |> Axon.dense(intermediate_size * 2, name: join(name, "intermediate"))
      |> Axon.split(2, axis: -1)

    x = Axon.multiply(x, Axon.silu(gate))

    x
    |> Axon.dropout(rate: dropout, name: join(name, "dropout"))
    |> Axon.dense(output_size, name: join(name, "output"))
  end

  # :norm_first block with additional scaling layers
  defp block_impl(hidden_state, steps, name, spec) do
    shortcut = hidden_state

    {hidden_state, attention_info} =
      hidden_state
      |> steps.self_attention_norm.()
      |> steps.self_attention.()

    hidden_state =
      hidden_state
      |> Bumblebee.Layers.scale(
        scale_initializer: Axon.Initializers.full(spec.scale_initial_value),
        name: join(name, "self_attention_scale")
      )
      |> Axon.add(shortcut)

    {_hidden_state, cross_attention_info} =
      steps.cross_attention_maybe.(hidden_state, fn _hidden_state ->
        raise "cross attention not supported"
      end)

    shortcut = hidden_state

    hidden_state =
      hidden_state
      |> steps.output_norm.()
      |> steps.ffn.()
      |> Bumblebee.Layers.scale(
        scale_initializer: Axon.Initializers.full(spec.scale_initial_value),
        name: join(name, "output_scale")
      )
      |> Axon.add(shortcut)

    {hidden_state, attention_info, cross_attention_info}
  end

  defp kernel_initializer(spec) do
    Axon.Initializers.normal(scale: spec.initializer_scale)
  end

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

      opts =
        convert!(data,
          image_size: {"image_size", number()},
          num_channels: {"num_channels", number()},
          patch_size: {"patch_size", number()},
          hidden_size: {"hidden_size", number()},
          num_blocks: {"num_hidden_layers", number()},
          num_attention_heads: {"num_attention_heads", number()},
          intermediate_size_ratio: {"mlp_ratio", number()},
          activation: {"hidden_act", activation()},
          use_attention_bias: {"qkv_bias", boolean()},
          dropout_rate: {"hidden_dropout_prob", number()},
          attention_dropout_rate: {"attention_probs_dropout_prob", number()},
          layer_norm_epsilon: {"layer_norm_eps", number()},
          initializer_scale: {"initializer_range", number()},
          scale_initial_value: {"layerscale_value", number()},
          ffn_swiglu_activation: {"use_swiglu_ffn", boolean()},
          backbone_output_indices: {"out_indices", list(number())},
          backbone_use_norm: {"use_backbone_norm", 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.patch_embedding.projection" => "dinov2.embeddings.patch_embeddings.projection",
        "embedder.class_embedding" => %{
          "embeddings" => {
            [{"dinov2.embeddings", "cls_token"}],
            fn [value] -> Nx.squeeze(value, axes: [0]) end
          }
        },
        "embedder.position_embedding" => %{
          "embeddings" => {
            [{"dinov2.embeddings", "position_embeddings"}],
            fn [value] -> Nx.squeeze(value, axes: [0]) end
          }
        },
        "encoder.blocks.{n}.self_attention_norm" => "dinov2.encoder.layer.{n}.norm1",
        "encoder.blocks.{n}.self_attention.key" =>
          "dinov2.encoder.layer.{n}.attention.attention.key",
        "encoder.blocks.{n}.self_attention.query" =>
          "dinov2.encoder.layer.{n}.attention.attention.query",
        "encoder.blocks.{n}.self_attention.value" =>
          "dinov2.encoder.layer.{n}.attention.attention.value",
        "encoder.blocks.{n}.self_attention.output" =>
          "dinov2.encoder.layer.{n}.attention.output.dense",
        "encoder.blocks.{n}.self_attention_scale" => %{
          "scale" => {
            [{"dinov2.encoder.layer.{n}.layer_scale1", "lambda1"}],
            fn [lambda1] -> lambda1 end
          }
        },
        "encoder.blocks.{n}.ffn.intermediate" =>
          if(spec.ffn_swiglu_activation,
            do: "dinov2.encoder.layer.{n}.mlp.weights_in",
            else: "dinov2.encoder.layer.{n}.mlp.fc1"
          ),
        "encoder.blocks.{n}.ffn.output" =>
          if(spec.ffn_swiglu_activation,
            do: "dinov2.encoder.layer.{n}.mlp.weights_out",
            else: "dinov2.encoder.layer.{n}.mlp.fc2"
          ),
        "encoder.blocks.{n}.output_norm" => "dinov2.encoder.layer.{n}.norm2",
        "encoder.blocks.{n}.output_scale" => %{
          "scale" => {
            [{"dinov2.encoder.layer.{n}.layer_scale2", "lambda1"}],
            fn [lambda1] -> lambda1 end
          }
        },
        "norm" => "dinov2.layernorm",
        "image_classification_head.output" => "classifier"
      }
    end
  end
end