"""Immerse detail — overlay high-resolution detail onto lower-resolution overview."""

from __future__ import annotations

import numpy as np

from backend.node_registry import register_node
from backend.data_types import DataField


@register_node(display_name="Immerse Detail")
class ImmerseDetail:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "overview": ("DATA_FIELD",),
                "detail": ("DATA_FIELD",),
                "blend": (["replace", "average"], {"default": "replace"}),
            }
        }

    OUTPUTS = (
        ('DATA_FIELD', 'combined'),
    )
    FUNCTION = "process"

    DESCRIPTION = (
        "Overlay a high-resolution detail scan onto a lower-resolution overview "
        "image using cross-correlation to find the best position. "
    )

    KEYWORDS = ("inset", "zoom", "merge", "overlay")

    def process(self, overview: DataField, detail: DataField, blend: str) -> tuple:
        ov = np.asarray(overview.data, dtype=np.float64)
        dt = np.asarray(detail.data, dtype=np.float64)

        # Resample detail to overview pixel size if needed
        scale_x = detail.dx / overview.dx
        scale_y = detail.dy / overview.dy

        if abs(scale_x - 1.0) > 0.01 or abs(scale_y - 1.0) > 0.01:
            from scipy.ndimage import zoom
            dt = zoom(dt, (scale_y, scale_x), order=1)

        dy_res, dx_res = dt.shape
        oy_res, ox_res = ov.shape

        if dy_res >= oy_res or dx_res >= ox_res:
            # Detail is larger than overview, just return overview
            return (overview,)

        # Cross-correlate to find best position
        # Use a sliding window approach for small detail
        best_score = -np.inf
        best_y, best_x = 0, 0

        dt_norm = dt - dt.mean()
        dt_std = dt.std()
        if dt_std < 1e-30:
            dt_std = 1.0

        # Coarse search with stride
        stride = max(1, min(dy_res, dx_res) // 4)
        for iy in range(0, oy_res - dy_res + 1, stride):
            for ix in range(0, ox_res - dx_res + 1, stride):
                patch = ov[iy:iy + dy_res, ix:ix + dx_res]
                score = np.sum((patch - patch.mean()) * dt_norm) / (patch.std() * dt_std + 1e-30)
                if score > best_score:
                    best_score = score
                    best_y, best_x = iy, ix

        # Fine search around best position
        for iy in range(max(0, best_y - stride), min(oy_res - dy_res + 1, best_y + stride + 1)):
            for ix in range(max(0, best_x - stride), min(ox_res - dx_res + 1, best_x + stride + 1)):
                patch = ov[iy:iy + dy_res, ix:ix + dx_res]
                score = np.sum((patch - patch.mean()) * dt_norm) / (patch.std() * dt_std + 1e-30)
                if score > best_score:
                    best_score = score
                    best_y, best_x = iy, ix

        # Place detail into overview
        result = ov.copy()
        if blend == "replace":
            result[best_y:best_y + dy_res, best_x:best_x + dx_res] = dt
        else:  # average
            result[best_y:best_y + dy_res, best_x:best_x + dx_res] = \
                0.5 * (ov[best_y:best_y + dy_res, best_x:best_x + dx_res] + dt)

        return (overview.replace(data=result),)