"""
Geometry Primitives Library

Geometric transforms: rotate, scale, flip, translate, remap.
"""
import numpy as np
import cv2


def prim_translate(img, dx, dy):
    """Translate image by (dx, dy) pixels."""
    h, w = img.shape[:2]
    M = np.float32([[1, 0, dx], [0, 1, dy]])
    return cv2.warpAffine(img, M, (w, h))


def prim_rotate(img, angle, cx=None, cy=None):
    """Rotate image by angle degrees around center (cx, cy)."""
    h, w = img.shape[:2]
    if cx is None:
        cx = w / 2
    if cy is None:
        cy = h / 2
    M = cv2.getRotationMatrix2D((cx, cy), angle, 1.0)
    return cv2.warpAffine(img, M, (w, h))


def prim_scale(img, sx, sy, cx=None, cy=None):
    """Scale image by (sx, sy) around center (cx, cy)."""
    h, w = img.shape[:2]
    if cx is None:
        cx = w / 2
    if cy is None:
        cy = h / 2

    # Build transform matrix
    M = np.float32([
        [sx, 0, cx * (1 - sx)],
        [0, sy, cy * (1 - sy)]
    ])
    return cv2.warpAffine(img, M, (w, h))


def prim_flip_h(img):
    """Flip image horizontally."""
    return cv2.flip(img, 1)


def prim_flip_v(img):
    """Flip image vertically."""
    return cv2.flip(img, 0)


def prim_flip(img, direction="horizontal"):
    """Flip image in given direction."""
    if direction in ("horizontal", "h"):
        return prim_flip_h(img)
    elif direction in ("vertical", "v"):
        return prim_flip_v(img)
    elif direction in ("both", "hv", "vh"):
        return cv2.flip(img, -1)
    return img


def prim_transpose(img):
    """Transpose image (swap x and y)."""
    return np.transpose(img, (1, 0, 2))


def prim_remap(img, map_x, map_y):
    """Remap image using coordinate maps."""
    return cv2.remap(img, map_x.astype(np.float32),
                     map_y.astype(np.float32),
                     cv2.INTER_LINEAR)


def prim_make_coords(w, h):
    """Create coordinate grids for remapping."""
    x = np.arange(w, dtype=np.float32)
    y = np.arange(h, dtype=np.float32)
    map_x, map_y = np.meshgrid(x, y)
    return (map_x, map_y)


def prim_perspective(img, src_pts, dst_pts):
    """Apply perspective transform."""
    src = np.float32(src_pts)
    dst = np.float32(dst_pts)
    M = cv2.getPerspectiveTransform(src, dst)
    h, w = img.shape[:2]
    return cv2.warpPerspective(img, M, (w, h))


def prim_affine(img, src_pts, dst_pts):
    """Apply affine transform using 3 point pairs."""
    src = np.float32(src_pts)
    dst = np.float32(dst_pts)
    M = cv2.getAffineTransform(src, dst)
    h, w = img.shape[:2]
    return cv2.warpAffine(img, M, (w, h))


def _get_legacy_geometry_primitives():
    """Import geometry primitives from legacy primitives module."""
    from sexp_effects.primitives import (
        prim_coords_x,
        prim_coords_y,
        prim_ripple_displace,
        prim_fisheye_displace,
        prim_kaleidoscope_displace,
    )
    return {
        'coords-x': prim_coords_x,
        'coords-y': prim_coords_y,
        'ripple-displace': prim_ripple_displace,
        'fisheye-displace': prim_fisheye_displace,
        'kaleidoscope-displace': prim_kaleidoscope_displace,
    }


PRIMITIVES = {
    # Basic transforms
    'translate': prim_translate,
    'rotate-img': prim_rotate,
    'scale-img': prim_scale,

    # Flips
    'flip-h': prim_flip_h,
    'flip-v': prim_flip_v,
    'flip': prim_flip,
    'transpose': prim_transpose,

    # Remapping
    'remap': prim_remap,
    'make-coords': prim_make_coords,

    # Advanced transforms
    'perspective': prim_perspective,
    'affine': prim_affine,

    # Displace / coordinate ops (from legacy primitives)
    **_get_legacy_geometry_primitives(),
}