"""Field translation (spatial shifting) utilities for CCD analysis.

When computing correction fields across scenes with different object positions
(e.g., illusion pinball at x=393 vs cloak pinball at x=613), fields must be
translated so that reference points align before subtraction.

All functions operate on field arrays with shape (NY, NX) or (N, NY, NX).
NX = 1280, NY = 512 in the standard grid.
"""
from __future__ import annotations

import numpy as np


def translate_field_x(field: np.ndarray, shift_x: int,
                      fill_edge: bool = True) -> np.ndarray:
    """Horizontally shift a field by a given number of pixels.

    Parameters
    ----------
    field : (NY, NX) or (N, NY, NX) ndarray
        Velocity field(s) in legacy (NY, NX) order.
    shift_x : int
        Positive = shift right, negative = shift left.
    fill_edge : bool
        If True, fill vacated columns with edge values (smooth).
        If False, fill with zeros (creates boundary artifacts).

    Returns
    -------
    shifted : ndarray with same shape as input.
    """
    if shift_x == 0:
        return field.copy()

    if field.ndim == 2:
        ny, nx = field.shape
        result = np.zeros_like(field)
        shift = shift_x
        if shift > 0:
            result[:, shift:] = field[:, :-shift]
            if fill_edge:
                # Fill left vacated columns with leftmost column value
                result[:, :shift] = field[:, :1]
        else:
            s = -shift
            result[:, :-s] = field[:, s:]
            if fill_edge:
                # Fill right vacated columns with rightmost column value
                result[:, -s:] = field[:, -1:]
        return result
    elif field.ndim == 3:
        n, ny, nx = field.shape
        result = np.zeros_like(field)
        shift = shift_x
        if shift > 0:
            result[:, :, shift:] = field[:, :, :-shift]
            if fill_edge:
                # Broadcast leftmost column across vacated columns
                result[:, :, :shift] = field[:, :, :1]
        else:
            s = -shift
            result[:, :, :-s] = field[:, :, s:]
            if fill_edge:
                result[:, :, -s:] = field[:, :, -1:]
        return result
    else:
        raise ValueError(f"Unsupported field ndim: {field.ndim}")


def translate_fields_dict(q: dict, shift_x: int) -> dict:
    """Translate ux and uy fields in a data dict, leaving telemetry unchanged."""
    if shift_x == 0:
        return q
    return {
        "ux": translate_field_x(q["ux"], shift_x),
        "uy": translate_field_x(q["uy"], shift_x),
        "forces": q.get("forces"),
        "sensors": q.get("sensors"),
        "actions": q.get("actions"),
        "meta": {**q.get("meta", {}), "translated_by": shift_x},
        "step_indices": q.get("step_indices"),
    }


# ---------------------------------------------------------------------------
# Reference positions (pixel coordinates)
# ---------------------------------------------------------------------------
# All scenes now use UNIFIED geometry:
#   pinball at (30, 31.3) x L0 -> center = 613 px
#   sensors at 40 x L0
#   target cylinder at 30.65 x L0 = 613 px (same as pinball center)
# field_translate kept for optional cross-comparison or future use.

ILLUSION_REF_X = 613    # UNIFIED: was 393
CLOAK_REF_X = 613       # unchanged
TARGET_CYL_REF_X = 613  # UNIFIED: was 400 (now same as pinball center)


def get_scene_ref_x(scene_name: str) -> int | None:
    """Get the reference x-position (pinball/cylinder center) for a scene.

    Returns pixel coordinate, or None if no reference (e.g. target_channel).
    """
    if "illusion" in scene_name:
        return ILLUSION_REF_X
    if "target_cylinder" in scene_name:
        return TARGET_CYL_REF_X
    if scene_name in ("pinball", "steady_cloak", "karman_re100", "karman_q_blk",
                      "vortex_lamb", "vortex_taylor",
                      "vortex_uncontrolled_lamb", "vortex_uncontrolled_taylor"):
        return CLOAK_REF_X
    return None
