"""
	"LiftFeat: 3D Geometry-Aware Local Feature Matching"
"""


import sys
import os

ALIKE_PATH = '/home/yepeng_liu/code_python/multimodal_remote/ALIKE'
sys.path.append(ALIKE_PATH)

import torch
import torch.nn as nn
from alike import ALike
import cv2
import numpy as np

import pdb

dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

configs = {
    'alike-t': {'c1': 8, 'c2': 16, 'c3': 32, 'c4': 64, 'dim': 64, 'single_head': True, 'radius': 2,
                'model_path': os.path.join(ALIKE_PATH, 'models', 'alike-t.pth')},
    'alike-s': {'c1': 8, 'c2': 16, 'c3': 48, 'c4': 96, 'dim': 96, 'single_head': True, 'radius': 2,
                'model_path': os.path.join(ALIKE_PATH, 'models', 'alike-s.pth')},
    'alike-n': {'c1': 16, 'c2': 32, 'c3': 64, 'c4': 128, 'dim': 128, 'single_head': True, 'radius': 2,
                'model_path': os.path.join(ALIKE_PATH, 'models', 'alike-n.pth')},
    'alike-l': {'c1': 32, 'c2': 64, 'c3': 128, 'c4': 128, 'dim': 128, 'single_head': False, 'radius': 2,
                'model_path': os.path.join(ALIKE_PATH, 'models', 'alike-l.pth')},
}


class ALikeExtractor(nn.Module):
    def __init__(self,model_type,device) -> None:
        super().__init__()
        self.net=ALike(**configs[model_type],device=device,top_k=4096,scores_th=0.1,n_limit=8000)
        
    
    @torch.inference_mode()
    def extract_alike_kpts(self,img):
        pred0=self.net(img,sub_pixel=True)
        return pred0['keypoints']