"""Plot functions for kriging."""

# Copyright 2024 Koji Noshita
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import numpy as np
import numpy.typing as npt

# NOTE: Intentional cross-boundary import from landmark's private module.
# plot._kriging is the only external consumer of these TPS functions.
# If _kernels is refactored, update this import accordingly.
from ..landmark._kernels import tps_coefficients, tps_warp
from ._base import require_dependencies


def tps_grid_2d_plot(
    x_reference: npt.ArrayLike,
    x_target: npt.ArrayLike,
    grid_size: str | float = "auto",
    outer: float = 0.1,
    n_grid_inner: int = 10,
    ax: object | None = None,
) -> object:
    """Plot the thin-plate spline 2D warped grid.

    Parameters
    ----------
    x_reference : array-like, shape (n_landmarks, n_dim)
        Reference configuration.
    x_target : array-like, shape (n_landmarks, n_dim)
        Target configuration.
    grid_size : str/float, optional
        Grid size, by default "auto"
    outer : float, optional
        Outer range of x_reference covered by the grid, by default 0.1
    n_grid_inner : int, optional
        Number of inner points on each grid, by default 10
    ax : :class:`matplotlib.axes.Axes`, optional
        Pre-existing axes for the plot. Otherwise, a new figure and axes
        are created.

    Returns
    -------
    ax : :class:`matplotlib.axes.Axes`
        Matplotlib axes.

    Raises
    ------
    ImportError
        If matplotlib is not installed.
    """
    require_dependencies("matplotlib")
    import matplotlib.pyplot as plt

    x_reference = np.asarray(x_reference)
    W, c, A = tps_coefficients(x_reference, x_target)

    if ax is None:
        _, ax = plt.subplots()

    ref_min = np.min(x_reference, axis=0)
    ref_max = np.max(x_reference, axis=0)
    span = ref_max - ref_min
    x_min, y_min = ref_min - outer * span
    x_max, y_max = ref_max + outer * span

    w = x_max - x_min
    h = y_max - y_min

    grid_size_ = grid_size
    if grid_size == "auto":
        grid_size_ = np.min([w, h]) / 10

    if w > h:
        w = w - w % grid_size_ + grid_size_
    else:
        h = h - h % grid_size_ + grid_size_
    n_grid_x = np.rint(w / grid_size_)
    n_grid_y = np.rint(h / grid_size_)

    n_grid_x_ = int(n_grid_x * n_grid_inner + 1)
    n_grid_y_ = int(n_grid_y * n_grid_inner + 1)

    # Grid points
    xx = np.linspace(x_min, x_max, n_grid_x_)
    yy = np.linspace(y_min, y_max, n_grid_y_)
    grid_x, grid_y = np.meshgrid(xx, yy, indexing="ij")
    grid_points = np.column_stack([grid_x.ravel(), grid_y.ravel()])

    # Warp all grid points
    warped = tps_warp(grid_points, x_reference, W, c, A)
    w_1 = warped.reshape(n_grid_x_, n_grid_y_, 2)
    w_2 = w_1.transpose(1, 0, 2)

    ax.plot(w_1[:, ::n_grid_inner, 0], w_1[:, ::n_grid_inner, 1], "gray")
    ax.plot(w_2[:, ::n_grid_inner, 0], w_2[:, ::n_grid_inner, 1], "gray")
    ax.axis("equal")

    ax.scatter(x=x_reference[:, 0], y=x_reference[:, 1], zorder=2)
    ax.scatter(x=x_target[:, 0], y=x_target[:, 1], zorder=2)
    return ax
