# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Wrapper for performing OmniGlue inference, plus (optionally) SP/DINO."""

import numpy as np
from omniglue import dino_extract
from omniglue import superpoint_extract
from omniglue import utils
import tensorflow as tf

DINO_FEATURE_DIM = 768
MATCH_THRESHOLD = 1e-3


class OmniGlue:
  # TODO(omniglue): class docstring

  def __init__(
      self,
      og_export: str,
      sp_export: str | None = None,
      dino_export: str | None = None,
  ) -> None:
    self.matcher = tf.saved_model.load(og_export)
    if sp_export is not None:
      self.sp_extract = superpoint_extract.SuperPointExtract(sp_export)
    if dino_export is not None:
      self.dino_extract = dino_extract.DINOExtract(dino_export, feature_layer=1)

  def FindMatches(self, image0: np.ndarray, image1: np.ndarray):
    """TODO(omniglue): docstring."""
    height0, width0 = image0.shape[:2]
    height1, width1 = image1.shape[:2]
    sp_features0 = self.sp_extract(image0)
    sp_features1 = self.sp_extract(image1)
    dino_features0 = self.dino_extract(image0)
    dino_features1 = self.dino_extract(image1)
    dino_descriptors0 = dino_extract.get_dino_descriptors(
        dino_features0,
        tf.convert_to_tensor(sp_features0[0], dtype=tf.float32),
        tf.convert_to_tensor(height0, dtype=tf.int32),
        tf.convert_to_tensor(width0, dtype=tf.int32),
        DINO_FEATURE_DIM,
    )
    dino_descriptors1 = dino_extract.get_dino_descriptors(
        dino_features1,
        tf.convert_to_tensor(sp_features1[0], dtype=tf.float32),
        tf.convert_to_tensor(height1, dtype=tf.int32),
        tf.convert_to_tensor(width1, dtype=tf.int32),
        DINO_FEATURE_DIM,
    )

    inputs = self._construct_inputs(
        width0,
        height0,
        width1,
        height1,
        sp_features0,
        sp_features1,
        dino_descriptors0,
        dino_descriptors1,
    )

    og_outputs = self.matcher.signatures['serving_default'](**inputs)
    soft_assignment = og_outputs['soft_assignment'][:, :-1, :-1]

    match_matrix = (
        utils.soft_assignment_to_match_matrix(soft_assignment, MATCH_THRESHOLD)
        .numpy()
        .squeeze()
    )

    # Filter out any matches with 0.0 confidence keypoints.
    match_indices = np.argwhere(match_matrix)
    keep = []
    for i in range(match_indices.shape[0]):
      match = match_indices[i, :]
      if (sp_features0[2][match[0]] > 0.0) and (
          sp_features1[2][match[1]] > 0.0
      ):
        keep.append(i)
    match_indices = match_indices[keep]

    # Format matches in terms of keypoint locations.
    match_kp0s = []
    match_kp1s = []
    match_confidences = []
    for match in match_indices:
      match_kp0s.append(sp_features0[0][match[0], :])
      match_kp1s.append(sp_features1[0][match[1], :])
      match_confidences.append(soft_assignment[0, match[0], match[1]])
    match_kp0s = np.array(match_kp0s)
    match_kp1s = np.array(match_kp1s)
    match_confidences = np.array(match_confidences)
    return match_kp0s, match_kp1s, match_confidences

  ### Private methods ###

  def _construct_inputs(
      self,
      width0,
      height0,
      width1,
      height1,
      sp_features0,
      sp_features1,
      dino_descriptors0,
      dino_descriptors1,
  ):
    inputs = {
        'keypoints0': tf.convert_to_tensor(
            np.expand_dims(sp_features0[0], axis=0),
            dtype=tf.float32,
        ),
        'keypoints1': tf.convert_to_tensor(
            np.expand_dims(sp_features1[0], axis=0), dtype=tf.float32
        ),
        'descriptors0': tf.convert_to_tensor(
            np.expand_dims(sp_features0[1], axis=0), dtype=tf.float32
        ),
        'descriptors1': tf.convert_to_tensor(
            np.expand_dims(sp_features1[1], axis=0), dtype=tf.float32
        ),
        'scores0': tf.convert_to_tensor(
            np.expand_dims(np.expand_dims(sp_features0[2], axis=0), axis=-1),
            dtype=tf.float32,
        ),
        'scores1': tf.convert_to_tensor(
            np.expand_dims(np.expand_dims(sp_features1[2], axis=0), axis=-1),
            dtype=tf.float32,
        ),
        'descriptors0_dino': tf.expand_dims(dino_descriptors0, axis=0),
        'descriptors1_dino': tf.expand_dims(dino_descriptors1, axis=0),
        'width0': tf.convert_to_tensor(
            np.expand_dims(width0, axis=0), dtype=tf.int32
        ),
        'width1': tf.convert_to_tensor(
            np.expand_dims(width1, axis=0), dtype=tf.int32
        ),
        'height0': tf.convert_to_tensor(
            np.expand_dims(height0, axis=0), dtype=tf.int32
        ),
        'height1': tf.convert_to_tensor(
            np.expand_dims(height1, axis=0), dtype=tf.int32
        ),
    }
    return inputs