import torch
import numpy as np

from .functions import *
from typing import Union


class PolyHarmInterpolator(torch.nn.Module):
  r"""
  Interpolate batched data using polyharmonic interpolation.

  The interpolant has the form:
  
  :math:`f(\mathbf{x}) = \sum_{i=1}^{n} w_i\, \varphi(\|\mathbf{x} - \mathbf{c}_i\|)
    + \mathbf{v}^\text{T} \begin{bmatrix} 1 \\ \mathbf{x} \end{bmatrix}`.

  This is a sum of two terms: (1) a weighted sum of radial basis function (RBF) 
  terms with centers :math:`{\left(\mathbf{c}_1, \ldots, \mathbf{c}_n\right)}` 
  and (2) a linear term with a bias. The :math:`{\mathbf{c}_i}` vectors are 
  'training' points. The coefficients :math:`\mathbf{w}` and :math:`\mathbf{v}` 
  are estimated such that the interpolant exactly fits the value of the function 
  at the :math:`{\mathbf{c}_i}` points, and the vector :math:`\mathbf{w}` is 
  orthogonal to each :math:`{\mathbf{c}_i}`, and the vector :math:`\mathbf{w}` 
  sums to 0. With these constraints, the coefficients can be obtained by solving 
  a linear system.

  The function :math:`{\varphi}` is an RBF, parametrized by an interpolation
  order. Using `order=2` produces the well-known thin-plate spline.

  We also provide the option to perform regularized interpolation. Here, the
  interpolant is selected to trade off between the squared loss on the
  training data and a certain measure of its curvature
  [`details <https://en.wikipedia.org/wiki/Polyharmonic_spline>`_].
  Using a regularization weight greater than zero has the effect that the
  interpolant will no longer exactly fit the training data. However, it may be
  less vulnerable to overfitting, particularly for high-order interpolation.

  :param c: 3D tensor with shape `[batch_size, n, d]` of `n` `d`-dimensional
            locations. These do not need to be regularly-spaced.
  :type c: torch.Tensor or np.ndarray
  :param f: 3D tensor with shape `[batch_size, n, k]` of `n` `k`-dimensional
            values evaluated at train_points.
  :type f: torch.Tensor or np.ndarray
  :param order: (optional) Order of the interpolation. Common values are
                1 for :math:`{\varphi(r)=r}`, 2 for :math:`{\varphi(r)=r^2 \log(r)}`
                (thin-plate spline), or 3 for :math:`{\varphi(r)=r^3}`.
  :type order: int
  :param smoothing: (optional) Weight placed on the regularization term.
                    This will depend substantially on the problem, and
                    it should always be tuned. For many problems, it is
                    reasonable to use no regularization. If using a non-zero
                    value, we recommend a small value like 0.001.
  :type smoothing: float
  :param device: (optional) Specifies the default device to store tensors
                 and perform interpolation.
  :type device: str
  :param dtype: (optional) Specifies the default precision.
  :type dtype: torch.dtype
  """
  def __init__(
    self,
    c: Union[torch.Tensor, np.ndarray],
    f: Union[torch.Tensor, np.ndarray],
    order: int = 3,
    smoothing: float = 0.0,
    device: str = "cpu",
    dtype: torch.dtype = torch.float,
    *args,
    **kwargs
  ) -> None:
    super(PolyHarmInterpolator, self).__init__(*args, **kwargs)
    # Set dtype and device
    self.factory_kwargs = {
      "dtype": dtype,
      "device": device
    }
    # Set training data
    for (k, x) in (('c', c), ('f', f)):
      if isinstance(x, np.ndarray):
        x = torch.from_numpy(x)
      if (x.ndim != 3):
        raise ValueError(f"'{k}' must be a 3-dimensional tensor.")
      x = x.to(**self.factory_kwargs)
      self.register_buffer(k, x)
    if (self.c.shape[:2] != self.f.shape[:2]):
      raise ValueError(
        "The first two dimensions of 'c' and 'f' must be the same."
      )
    # Get smoothing and kernel function
    self.smoothing = float(smoothing)
    self.order = int(order)
    self.phi = get_phi(self.order)
    # Fit the interpolant to the observed data
    self._built = False
    self.build()

  def build(self) -> None:
    r"""
    Solve for interpolation coefficients.

    Computes the coefficients :math:`\mathbf{w}` and :math:`\mathbf{v}` of the
    polyharmonic interpolant for the training data using the kernel :math:`{\varphi}`.
    """
    # Get dimensions
    b, n, d = list(self.c.shape)
    k = self.f.shape[-1]
    # Construct the linear system
    # > Matrix A
    amat = self.phi(pairwise_squared_distance_matrix(self.c))
    if (self.smoothing > 0):
      imat = torch.unsqueeze(torch.eye(n, **self.factory_kwargs), dim=0)
      amat += self.smoothing * imat
    # > Matrix B
    ones = torch.ones_like(self.c[..., :1], **self.factory_kwargs)
    bmat = torch.cat([self.c, ones], dim=2)
    bmat_ncols = bmat.shape[2]
    # > Left hand side
    lhs_zeros = torch.zeros([b, bmat_ncols, bmat_ncols], **self.factory_kwargs)
    block_right = torch.cat([bmat, lhs_zeros], dim=1)
    block_left = torch.cat([amat, torch.permute(bmat, dims=(0,2,1))], dim=1)
    lhs = torch.cat([block_left, block_right], dim=2)
    # > Right hand side
    rhs_zeros = torch.zeros([b, d + 1, k], **self.factory_kwargs)
    rhs = torch.cat([self.f, rhs_zeros], dim=1)
    # Solve the linear system
    w_v = torch.linalg.solve(lhs, rhs)
    w, v = w_v[:, :n, :], w_v[:, n:, :]
    self.register_buffer("w", w)
    self.register_buffer("v", v)
    self._built = True

  def forward(
    self,
    x: Union[torch.Tensor, np.ndarray]
  ) -> torch.Tensor:
    r"""
    Apply polyharmonic interpolation model to new input data.

    Given coefficients :math:`\mathbf{w}` and :math:`\mathbf{v}` for the interpolation
    model, the interpolated function is evaluated at query points :math:`{\mathbf{x}}`.

    Note that the interpolation procedure is differentiable with respect
    to :math:`{x}`.

    :param x: 3D tensor with shape `[batch_size, m, d]`
              to evaluate the interpolation at.
    :type x: torch.Tensor

    :return: Polyharmonic interpolation evaluated at query points `x`.
    :rtype: torch.Tensor

    :raises ValueError: If the input tensor `x` is not 3-dimensional.
    """
    if (not self._built):
      raise ValueError("The interpolator has not been built.")
    if (x.ndim != 3):
      raise ValueError("'x' must be a 3-dimensional tensor.")
    if isinstance(x, np.ndarray):
      x = torch.from_numpy(x).to(**self.factory_kwargs)
    if (x.device != self.factory_kwargs["device"]):
      x = x.to(**self.factory_kwargs)
    # Compute the contribution from the rbf term
    d = cross_squared_distance_matrix(x, self.c)
    d_phi = self.phi(d)
    rbf_term = torch.matmul(d_phi, self.w)
    # Compute the contribution from the linear term
    ones = torch.ones_like(x[..., :1], **self.factory_kwargs)
    x_pad = torch.concat([x, ones], dim=2)
    linear_term = torch.matmul(x_pad, self.v)
    return rbf_term + linear_term
