import torch
import numpy as np
from PIL import Image

from unicorrn.model import build_model
from unicorrn.utils import safe_load_weights, plot_correspondences
from unicorrn.utils.config import read_yaml_config
from unicorrn.inference import (
    init_query_points,
    coarse_to_fine,
    cycle_uniform_grid_inference,
)

MODEL_CONFIG_PATH = "/your_project_path/configs/models/unicorrn_large_stage2.yml"
CKPT_PATH = "/your_project_path/pretrained_models/UniCorrn_Large_Stage2.pth"
GRID_SIZE = 4
CONFIDENCE_THRESHOLD = 3.8
MATCHING_RADIUS_PX = 1.0

# ---- Shared setup: load model and input images ----
model_cfg = read_yaml_config(MODEL_CONFIG_PATH)
model = build_model(model_cfg.NAME, cfg=model_cfg)
weights = torch.load(CKPT_PATH, map_location="cpu", weights_only=False)
safe_load_weights(model, weights["model"])
model = model.eval().cuda()

img1_path = "assets/image_a.png"
img2_path = "assets/image_b.png"

img1 = np.array(Image.open(img1_path).convert("RGB"))
img2 = np.array(Image.open(img2_path).convert("RGB"))

H, W = img1.shape[:2]

# ---- Usage 1: User-specified keypoints with confidence filter ----
queries = init_query_points(H, W, grid_size=GRID_SIZE).view(-1, 2).numpy()

kpts1, kpts2, confidence, _ = coarse_to_fine(
    img1,
    img2,
    queries,
    model,
    unified_model=True,
)

mask = confidence.squeeze() >= CONFIDENCE_THRESHOLD
kpts1 = kpts1[mask]
kpts2 = kpts2[mask]

plot_correspondences(img1, img2, kpts1, kpts2, marker_size=1, plot_line=False, save_path="example_confidence.jpg")

# ---- Usage 2: Cycle-consistency correspondence extraction ----
# Builds a uniform grid on img1 and keeps only matches whose backward
# prediction lands within `MATCHING_RADIUS_PX` of the original query.
cycle_kpts1, cycle_kpts2, _, _ = cycle_uniform_grid_inference(
    img1,
    img2,
    model,
    grid_size=GRID_SIZE,
    matching_radius_px=MATCHING_RADIUS_PX,
    unified_model=True,
)

plot_correspondences(img1, img2, cycle_kpts1, cycle_kpts2, marker_size=1, plot_line=False, save_path="example_cycle.jpg")