"""Dense matching stage (RoMa v2 for full and sparse modes)."""

import logging
from pathlib import Path

import torch

from ...dense import roma_warps_to_depth_maps, save_depth_map
from ...features import match_all_pairs_roma, run_roma_all_pairs
from ..context import PipelineContext
from ..helpers import _should_viz

logger = logging.getLogger(__name__)


def run_roma_full_path(
    undistorted_tensors: dict[str, torch.Tensor],
    ctx: PipelineContext,
    frame_dir: Path,
    frame_idx: int,
) -> tuple[dict, dict]:
    """Run RoMa v2 dense matching in full mode (warps to depth maps).

    Args:
        undistorted_tensors: Dict of undistorted BGR images as torch.Tensor.
        ctx: Pipeline context.
        frame_dir: Frame output directory.
        frame_idx: Frame index (for logging).

    Returns:
        Tuple of (depth_maps, confidence_maps):
            - depth_maps: Dict mapping camera name to depth map tensor.
            - confidence_maps: Dict mapping camera name to confidence map tensor.
    """
    config = ctx.config
    device = ctx.device

    # --- RoMa + full: warps -> depth maps -> fusion -> surface ---
    logger.info("Frame %d: running RoMa v2 dense matching (full mode)", frame_idx)
    all_warps = run_roma_all_pairs(
        undistorted_images=undistorted_tensors,
        pairs=ctx.pairs,
        config=config.dense_matching,
        device=device,
        masks=None,  # Masks applied after upsampling
    )

    logger.info("Frame %d: converting RoMa warps to depth maps", frame_idx)
    depth_maps, confidence_maps = roma_warps_to_depth_maps(
        ring_cameras=ctx.ring_cameras,
        pairs=ctx.pairs,
        all_warps=all_warps,
        projection_models=ctx.projection_models,
        dense_matching_config=config.dense_matching,
        reconstruction_config=config.reconstruction,
        image_size=list(ctx.calibration.cameras.values())[0].image_size,
        masks=ctx.masks,
    )

    # Save depth maps (opt-out)
    if config.runtime.save_depth_maps:
        depth_dir = frame_dir / "depth_maps"
        depth_dir.mkdir(exist_ok=True)
        for cam_name in depth_maps:
            save_depth_map(
                depth_maps[cam_name],
                confidence_maps[cam_name],
                depth_dir / f"{cam_name}.npz",
            )

    # [viz] Depth map renders
    if _should_viz(config, "depth"):
        try:
            from ...visualization.depth import render_all_depth_maps

            logger.info("Frame %d: rendering depth map visualizations", frame_idx)
            viz_dir = frame_dir / "viz"
            viz_dir.mkdir(exist_ok=True)

            # Convert depth/confidence tensors to numpy
            np_depths = {name: dm.cpu().numpy() for name, dm in depth_maps.items()}
            np_confs = {name: cm.cpu().numpy() for name, cm in confidence_maps.items()}

            render_all_depth_maps(np_depths, np_confs, viz_dir)
        except Exception:
            logger.exception("Frame %d: depth visualization failed", frame_idx)

    return depth_maps, confidence_maps


def run_roma_sparse_path(
    undistorted_tensors: dict[str, torch.Tensor],
    ctx: PipelineContext,
    frame_dir: Path,
    frame_idx: int,
) -> dict:
    """Run RoMa v2 in sparse mode (correspondences only).

    Args:
        undistorted_tensors: Dict of undistorted BGR images as torch.Tensor.
        ctx: Pipeline context.
        frame_dir: Frame output directory.
        frame_idx: Frame index (for logging).

    Returns:
        Dict with "all_matches" key containing match results.
    """
    config = ctx.config
    device = ctx.device

    # --- RoMa + sparse: correspondences -> triangulation -> sparse surface ---
    logger.info("Frame %d: matching with RoMa v2 (sparse mode)", frame_idx)
    all_matches = match_all_pairs_roma(
        undistorted_images=undistorted_tensors,
        pairs=ctx.pairs,
        config=config.dense_matching,
        device=device,
        masks=ctx.masks,
    )

    # Save matches if requested
    if config.runtime.save_features:
        from ...features import save_matches

        features_dir = frame_dir / "features"
        features_dir.mkdir(exist_ok=True)
        for (ref, src), match in all_matches.items():
            save_matches(match, features_dir / f"{ref}_{src}.pt")

    # Feature viz is skipped (no per-camera keypoints)

    return {"all_matches": all_matches}
