"""Conversion between pointwise and functional maps."""

import abc

import geomstats.backend as gs
import scipy
import torch
import torch.nn as nn
from sklearn.neighbors import NearestNeighbors

import geomfum.backend as xgs
import geomfum.wrap as _wrap  # noqa (for register)
from geomfum._registry import NeighborFinderRegistry, WhichRegistryMixins
from geomfum.neural_adjoint_map import NeuralAdjointMap


class BaseNeighborFinder(abc.ABC):
    """Base class for a Neighbor finder.

    Parameters
    ----------
    n_neighbors : int
        Number of neighbors to find.
    """

    def __init__(self, n_neighbors=1):
        self.n_neighbors = n_neighbors


class NeighborFinder(WhichRegistryMixins, BaseNeighborFinder):
    """Base class for a Neighbor finder.

    A simplified blueprint of ``sklearn.NearestNeighbors`` implementation.

    Parameters
    ----------
    n_neighbors : int
        Number of neighbors.
    """

    _Registry = NeighborFinderRegistry

    def __init__(self, n_neighbors=1):
        self.n_neighbors = n_neighbors
        self.sklearn_neighbor_finder = NearestNeighbors(
            n_neighbors=self.n_neighbors, leaf_size=40, algorithm="kd_tree", n_jobs=1
        )

    def __call__(self, X, Y):
        """Return indices of the points in `X` nearest to the points in `Y`.

        Parameters
        ----------
        X : array-like, shape=[n_points_x, n_features]
            Reference points.
        Y : array-like, shape=[n_points_y, n_features]
            Query points.

        Returns
        -------
        neigs : array-like, shape=[n_points_x, n_neighbors]
            Indices of the nearest neighbors in Y for each point in X.
        """
        self.sklearn_neighbor_finder.fit(xgs.to_device(Y, "cpu"))
        neigs = self.sklearn_neighbor_finder.kneighbors(
            xgs.to_device(X, "cpu"), return_distance=False
        )

        return gs.from_numpy(neigs)


class BaseP2pFromFmConverter(abc.ABC):
    """Pointwise map from functional map."""


class P2pFromFmConverter(BaseP2pFromFmConverter):
    """Pointwise map from functional map.

    Parameters
    ----------
    neighbor_finder : NeighborFinder
        Nearest neighbor finder.
    adjoint : bool
        Whether to use adjoint method.

    References
    ----------
    .. [OCSBG2012] Maks Ovsjanikov, Mirela Ben-Chen, Justin Solomon,
        Adrian Butscher, and Leonidas Guibas.
        “Functional Maps: A Flexible Representation of Maps between
        Shapes.” ACM Transactions on Graphics 31, no. 4 (2012): 30:1-30:11.
        https://doi.org/10.1145/2185520.2185526.
    .. [VM2023] Giulio Viganò  Simone Melzi. “Adjoint Bijective ZoomOut:
        Efficient Upsampling for Learned Linearly-Invariant Embedding.”
        The Eurographics Association, 2023. https://doi.org/10.2312/stag.20231293.
    """

    def __init__(self, neighbor_finder=None, adjoint=False, bijective=False):
        if neighbor_finder is None:
            neighbor_finder = NeighborFinder(n_neighbors=1)
        if neighbor_finder.n_neighbors > 1:
            raise ValueError("Expects `n_neighors = 1`.")

        self.neighbor_finder = neighbor_finder
        self.adjoint = adjoint
        self.bijective = bijective

    def __call__(self, fmap_matrix, basis_a, basis_b):
        """Convert functional map.

        Parameters
        ----------
        fmap_matrix : array-like, shape=[spectrum_size_b, spectrum_size_a]
            Functional map matrix.
        basis_a : Basis,
            Basis of the source shape.
        basis_b : Basis,
            Basis of the target shape.

        Returns
        -------
        p2p : array-like, shape=[{n_vertices_b, n_vertices_a}]
            Pointwise map. ``bijective`` controls shape.
        """
        k2, k1 = fmap_matrix.shape

        if self.adjoint:
            emb1 = basis_a.full_vecs[:, :k1]
            emb2 = basis_b.full_vecs[:, :k2] @ fmap_matrix

        else:
            emb1 = basis_a.full_vecs[:, :k1] @ fmap_matrix.T
            emb2 = basis_b.full_vecs[:, :k2]

        if self.bijective:
            emb1, emb2 = emb2, emb1

        p2p = self.neighbor_finder(emb2, emb1).flatten()

        return p2p


class SoftmaxNeighborFinder(BaseNeighborFinder, nn.Module):
    """Softmax neighbor finder.

    Finds neighbors using softmax regularization.

    Parameters
    ----------
    n_neighbors : int
        Number of neighbors.
    tau : float
        Temperature parameter for softmax regularization.
    """

    def __init__(self, n_neighbors=1, tau=0.07):
        BaseNeighborFinder.__init__(self, n_neighbors=n_neighbors)
        nn.Module.__init__(self)
        self.tau = tau

    def __call__(self, X, Y):
        """Return indices of the points in `X` nearest to the points in `Y`.

        Parameters
        ----------
        X : array-like, shape=[n_points_x, n_features]
            Reference points.
        Y : array-like, shape=[n_points_y, n_features]
            Query points.

        Returns
        -------
        neigs : array-like, shape=[n_points_x, n_neighbors]
            Indices of the nearest neighbors in Y for each point in X.
        """
        return self.forward(X, Y)

    def forward(self, X, Y):
        """Find k nearest neighbors using softmax regularization.

        Parameters
        ----------
        X : array-like, shape=[n_points_x, n_features]
            Reference points.
        Y : array-like, shape=[n_points_y, n_features]
            Query points.

        Returns
        -------
        neigs : array-like, shape=[n_points_x, n_neighbors]
            Indices of the nearest neighbors in Y for each point in X.
        """
        P = self.softmax_matrix(X, Y)
        # Get the indices of the top-k (self.n_neighbors) highest values for each row
        indices = torch.topk(P, self.n_neighbors, dim=-1)[1]
        return indices

    def softmax_matrix(self, X, Y):
        """Compute the permutation matrix P as a softmax of the similarity.

        Parameters
        ----------
        X : array-like, shape=[n_points_x, n_features]
            Reference points.
        Y : array-like, shape=[n_points_y, n_features]
            Query points.

        Returns
        -------
        P : array-like, shape=[n_points_x, n_points_y]
            Permutation matrix, where each row sums to 1.
        """
        similarity = torch.mm( X, Y.T)

        P = torch.exp(
            similarity / self.tau
            - torch.logsumexp(similarity / self.tau, dim=-1, keepdim=True)
        )

        return P


class SinkhornP2pFromFmConverter(P2pFromFmConverter):
    """Pointwise map from functional map using Sinkhorn filters.

    Parameters
    ----------
    neighbor_finder : SinkhornKNeighborsFinder
        Nearest neighbor finder.
    adjoint : bool
        Whether to use adjoint method.
    bijective : bool
        Whether to use bijective method. Check [VM2023]_.

    References
    ----------
    .. [PRMWO2021] Gautam Pai, Jing Ren, Simone Melzi, Peter Wonka, and Maks Ovsjanikov.
        "Fast Sinkhorn Filters: Using Matrix Scaling for Non-Rigid Shape Correspondence
        with Functional Maps." Proceedings of the IEEE/CVF Conference on Computer Vision
        and Pattern Recognition (CVPR), 2021, pp. 11956-11965.
        https://hal.science/hal-03184936/document
    """

    def __init__(
        self,
        neighbor_finder=None,
        adjoint=False,
        bijective=False,
    ):
        if neighbor_finder is None:
            neighbor_finder = NeighborFinder.from_registry(which="pot")

        super().__init__(
            neighbor_finder=neighbor_finder,
            adjoint=adjoint,
            bijective=bijective,
        )


class BaseFmFromP2pConverter(abc.ABC):
    """Functional map from pointwise map."""


class FmFromP2pConverter(BaseFmFromP2pConverter):
    """Functional map from pointwise map.

    Parameters
    ----------
    pseudo_inverse : bool
        Whether to solve using pseudo-inverse.
    """

    # TODO: add subsampling
    def __init__(self, pseudo_inverse=False):
        self.pseudo_inverse = pseudo_inverse

    def __call__(self, p2p, basis_a, basis_b):
        """Convert point to point map.

        Parameters
        ----------
        p2p : array-like, shape=[n_vertices_b]
            Pointwise map.
        basis_a : Basis,
            Basis of the source shape.
        basis_b : Basis,
            Basis of the target shape.

        Returns
        -------
        fmap_matrix : array-like, shape=[spectrum_size_b, spectrum_size_a]
            Functional map matrix.
        """
        evects1_pb = basis_a.vecs[p2p, :]

        if self.pseudo_inverse:
            return basis_b.vecs.T @ (basis_b._shape.laplacian.mass_matrix @ evects1_pb)

        return gs.from_numpy(scipy.linalg.lstsq(basis_b.vecs, evects1_pb)[0])


class FmFromP2pBijectiveConverter(BaseFmFromP2pConverter):
    """Bijective functional map from pointwise map method.

    References
    ----------
    .. [VM2023] Giulio Viganò  Simone Melzi. “Adjoint Bijective ZoomOut:
        Efficient Upsampling for Learned Linearly-Invariant Embedding.”
        The Eurographics Association, 2023. https://doi.org/10.2312/stag.20231293.
    """

    def __init__(self, pseudo_inverse=False):
        self.pseudo_inverse = pseudo_inverse

    def __call__(self, p2p, basis_a, basis_b):
        """Convert point to point map.

        Parameters
        ----------
        p2p : array-like, shape=[n_vertices_a]
            Pointwise map.
        basis_a : Basis,
            Basis of the source shape.
        basis_b : Basis,
            Basis of the target shape.

        Returns
        -------
        fmap_matrix : array-like, shape=[spectrum_size_b, spectrum_size_a]
            Functional map matrix.
        """
        evects2_pb = basis_b.vecs[p2p, :]

        if self.pseudo_inverse:
            return gs.linalg.pinv(evects2_pb) @ basis_a.vecs

        return gs.from_numpy(scipy.linalg.lstsq(evects2_pb, basis_a.vecs)[0])


class NamFromP2pConverter(BaseFmFromP2pConverter):
    """Neural Adjoint Map from pointwise map using Neural Adjoint Maps (NAMs)."""

    def __init__(self, iter_max=200, patience=10, min_delta=1e-4, device="cpu"):
        """Initialize the converter.

        Parameters
        ----------
        iter_max : int, optional
            Maximum number of iterations for training the Neural Adjoint Map.
        patience : int, optional
            Number of iterations with no improvement after which training will be stopped.
        min_delta : float, optional
            Minimum change in the loss to qualify as an improvement.
        device : str, optional
            Device to use for the Neural Adjoint Map (e.g., 'cpu' or 'cuda').
        """
        self.iter_max = iter_max
        self.device = device
        self.min_delta = min_delta
        self.patience = patience

    def __call__(self, p2p, basis_a, basis_b, optimizer=None):
        """Convert point to point map.

        Parameters
        ----------
        p2p : array-like, shape=[n_vertices_b]
            Pointwise map.
        basis_a : Basis,
            Basis of the source shape.
        basis_b : Basis,
            Basis of the target shape.
        optimizer : torch.optim.Optimizer, optional
            Optimizer for training the Neural Adjoint Map.

        Returns
        -------
        nam: NeuralAdjointMap , shape=[spectrum_size_b, spectrum_size_a]
            Neural Adjoint Map model.
        """
        evects1_pb = xgs.to_torch(basis_a.vecs[p2p, :]).to(self.device).double()
        evects2 = xgs.to_torch(basis_b.vecs).to(self.device).double()
        nam = NeuralAdjointMap(
            input_dim=basis_a.spectrum_size,
            output_dim=basis_b.spectrum_size,
            device=self.device,
        ).double()

        if optimizer is None:
            optimizer = torch.optim.Adam(nam.parameters(), lr=0.01, weight_decay=1e-5)

        best_loss = float("inf")
        wait = 0

        for _ in range(self.iter_max):
            optimizer.zero_grad()

            pred = nam(evects1_pb)

            loss = torch.nn.functional.mse_loss(pred, evects2)
            loss.backward()
            optimizer.step()

            if loss.item() < best_loss - self.min_delta:
                best_loss = loss.item()
                wait = 0
            else:
                wait += 1
            if wait >= self.patience:
                break

        return nam


class P2pFromNamConverter(BaseP2pFromFmConverter):
    """Pointwise map from Neural Adjoint Map (NAM).

    Parameters
    ----------
    neighbor_finder : NeighborFinder
        Nearest neighbor finder.
    """

    def __init__(self, neighbor_finder=None):
        if neighbor_finder is None:
            neighbor_finder = NeighborFinder(n_neighbors=1)
        if neighbor_finder.n_neighbors > 1:
            raise ValueError("Expects `n_neighors = 1`.")

        self.neighbor_finder = neighbor_finder

    def __call__(self, nam, basis_a, basis_b):
        """Convert neural adjoint map.

        Parameters
        ----------
        nam : NeuralAdjointMap, shape=[spectrum_size_b, spectrum_size_a]
            Nam model.
        basis_a : Basis,
            Basis of the source shape.
        basis_b : Basis,
            Basis of the target shape.

        Returns
        -------
        p2p : array-like, shape=[n_vertices_b]
            Pointwise map.
        """
        k2, k1 = nam.shape

        emb1 = nam(xgs.to_torch(basis_a.full_vecs[:, :k2]).to(nam.device).double())
        emb2 = xgs.to_torch(basis_b.full_vecs[:, :k1]).to(nam.device).double()

        p2p = self.neighbor_finder(emb2.detach().cpu(), emb1.detach().cpu()).flatten()
        return p2p
