# Copyright (c) 2026 The neuraLQX Authors - All rights reserved.
#
# 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.

r"""
Public graph interface built on top of :class:`neuralqx.graph._graph_handler.GraphHandler`.

This module defines :class:`AbstractGraph`, a thin, user facing wrapper that exposes a consistent API
for graphs used throughout neuraLQX. The underlying :class:`GraphHandler` is responsible for

- validating and normalising edge data, including multiedge keys
- constructing primal and dual graphs in both NetworkX and NetKet formats
- maintaining embeddings and coordinate maps for non planar graphs
- computing local sign data used in volume like and grasping operators

Planar versus non planar
In neuraLQX a planar graph uses integer vertex labels. A non planar graph is represented by edges
whose endpoints are 3 tuples :math:`(x,y,z)` giving an embedding in :math:`\mathbb R^3`. When
requested, a random embedding can be generated by reassigning each vertex coordinate from a normal
distribution

.. math::

   (x,y,z)\sim \mathcal N(\mu,\sigma^2 I_3)

The abstract interface provides hooks for minimal loop extraction and loop dressing, which are used
when constructing graph based operators such as plaquette terms and graph changing moves.
"""

import abc
import logging
from abc import abstractmethod
from typing import Optional
from typing import Union

from .graph_handler import GraphHandler

from neuralqx.debug import event
from neuralqx.utils.errors import InvalidIndexError
from neuralqx.utils.errors import NonExistentNonPlanarEdgesError
from neuralqx.utils.errors import NonExistentNonPlanarVerticesError
from neuralqx.utils.errors import OutOfRangeIndexError
from neuralqx.utils.errors import RandomEmbeddingForPlanarGraphWarning


def _is_really_planar(edges):
    """Checks if edges are really planar."""
    for edge in edges:
        if not isinstance(edge[0], tuple) or not isinstance(edge[1], tuple):
            return True
    return False


class AbstractGraph(abc.ABC):
    r"""
    Abstract base class for neuraLQX graphs.

    An :class:`AbstractGraph` owns a :class:`~neuralqx.graph._graph_handler.GraphHandler` instance and
    delegates all low level construction to it. The class exposes a stable API used by Hilbert spaces,
    operators, solvers, and callbacks, independent of whether the graph is planar or non planar.

    Edges
    -----
    Edges are treated as oriented multigraph edges. Each edge is represented as a triple
    ``(u, v, key)`` where ``u`` is the tail, ``v`` is the head, and ``key`` distinguishes parallel
    edges between the same endpoints.

    Dual graph
    ----------
    The dual graph is defined as the line graph of the primal graph, so dual vertices correspond to
    primal edges. The handler maintains a bijection between a primal edge and its dual vertex index.

    Minimal loops
    ---------------
    Concrete subclasses must implement access to minimal loops and their dressed variants. In this
    context, minimal loops refer to the shortest cycles in a cycle basis of the underlying undirected
    graph. Dressing attaches an orientation dependent label to each edge along the loop, which is
    typically used to select creation or annihilation type actions in operator kernels.

    Sign data
    ----------
    The property :pyattr:`signs` exposes local orientation signs used in volume like constructions. For
    non planar embeddings this is based on the determinant of three tangent vectors, while for planar
    embeddings this reduces to a 2D oriented area convention.

    :param edges: List of oriented edges. Each edge is ``(u, v)`` or ``(u, v, key)``.
    :param plot: If True, the handler may create and display diagnostic plots during construction.
    :param random_embedding: If True, treat the input as non planar and redraw all vertex coordinates
      from a Gaussian distribution.
    :param random_embedding_mean: Mean :math:`\mu` used when `random_embedding` is True.
    :param random_embedding_std: Standard deviation :math:`\sigma` used when `random_embedding` is True.
    :param random_embedding_seed: Seed controlling the random embedding when `random_embedding` is True.
    :raises ValueError: If the graph cannot be constructed from the provided edges.
    """

    handler: GraphHandler
    """A core handler instance."""

    def __init__(
        self,
        edges: list,
        plot: bool = False,
        *,
        random_embedding: bool = False,
        random_embedding_mean: float = 0.0,
        random_embedding_std: float = 5.0,
        random_embedding_seed: int = 123,
    ):
        r"""
        Construct a graph wrapper around a :class:`~neuralqx.graph._graph_handler.GraphHandler`.

        If `random_embedding` is enabled, vertex coordinates are reassigned by sampling new coordinates for
        each unique vertex label and rebuilding the edge list using the new coordinates. This is useful for
        generating generic non planar embeddings when only combinatorial connectivity is supplied.

        :param edges: List of oriented edges. Each edge is ``(u, v)`` or ``(u, v, key)``.
        :param plot: If True, enable plotting during handler initialisation.
        :param random_embedding: If True, randomise non planar coordinates before building the handler.
        :param random_embedding_mean: Mean of the Gaussian used to randomise coordinates.
        :param random_embedding_std: Standard deviation of the Gaussian used to randomise coordinates.
        :param random_embedding_seed: Seed for coordinate randomisation.
        :return: None.
        """

        event(
            msg="INITIALIZING GRAPH",
            tag="GRAPH:INIT",
            level=logging.INFO,
            edges=edges,
            plot=plot,
            random_embedding=random_embedding,
            random_embedding_mean=random_embedding_mean,
            random_embedding_std=random_embedding_std,
            random_embedding_seed=random_embedding_seed,
            graph_type=type(self).__name__,
        )

        # placeholder for the coordinates map
        coordinates = None

        # one or the other, embedding is only available for non-planar graphs
        if random_embedding and _is_really_planar(edges):
            RandomEmbeddingForPlanarGraphWarning()
            random_embedding = False

        if random_embedding:
            from .utils.embedding import randomize_graph_coordinates

            edges, coordinates = randomize_graph_coordinates(
                edges,
                random_embedding_mean,
                random_embedding_std,
                random_embedding_seed,
            )

        self.handler = GraphHandler(edges, coordinates, plot)

        self._hash = None

    @abstractmethod
    def minimal_loops(self):
        r"""
        Return minimal loops of the primal graph.

        A minimal loop is a shortest cycle selected from a cycle basis of the underlying undirected graph.
        Equivalently, if :math:`\mathcal C` denotes the cycle space and :math:`\mathcal B\subset\mathcal C`
        a chosen basis, this method returns the subset of cycles in :math:`\mathcal B` with minimal length.

        Concrete subclasses decide the precise representation, for example a list of edge triples
        ``(u, v, key)`` or a list of vertex sequences.

        :return: Minimal loops of the primal graph.
        :raises NotImplementedError: If the concrete graph class does not implement loop extraction.
        """

        raise NotImplementedError

    @abstractmethod
    def minimal_dual_loops(self, edge_rep: Optional[bool]):
        r"""
        Return minimal loops of the dual graph.

        The dual graph is the line graph of the primal graph. Dual vertices correspond to primal edges, so
        a dual loop corresponds to a sequence of primal edges with consecutive incidence.

        :param edge_rep: If True, return dual loops using primal edge representations. If False, return dual
          loops using dual vertex indices, depending on the concrete implementation.
        :return: Minimal loops of the dual graph.
        :raises NotImplementedError: If the concrete graph class does not implement loop extraction.
        """

        raise NotImplementedError

    @abstractmethod
    def dressed_minimal_loops(self):
        r"""
        Return dressed minimal loops of the primal graph.

        Dressing adds orientation dependent structure to each edge in the loop. A common convention is to
        associate to each oriented edge either a creation like label or an annihilation like label, which
        in operator language corresponds to selecting actions analogous to :math:`a^\dagger` or :math:`a`
        depending on whether the loop traversal agrees with the edge orientation.

        Concrete subclasses must specify the exact dressing schema and output format.

        :return: Dressed minimal loops of the primal graph.
        :raises NotImplementedError: If the concrete graph class does not implement dressed loop extraction.
        """

        raise NotImplementedError

    @abstractmethod
    def dressed_minimal_dual_loops(self):
        r"""
        Return dressed minimal loops of the dual graph.

        This is the dual analogue of :meth:`dressed_minimal_loops`. Dressing is applied after selecting the
        minimal loops in the dual graph and may be expressed in terms of primal edge orientation data.

        :return: Dressed minimal loops of the dual graph.
        :raises NotImplementedError: If the concrete graph class does not implement dressed dual loops.
        """

        raise NotImplementedError

    @property
    @abstractmethod
    def nk_graph(self):
        r"""
        Return the NetKet graph associated with the primal graph.

        Concrete subclasses may return a :class:`netket.graph.Graph` instance constructed from the handler.

        :return: NetKet primal graph.
        :raises NotImplementedError: If not implemented by the concrete graph class.
        """

        raise NotImplementedError

    @property
    @abstractmethod
    def dual_nk_graph(self):
        r"""
        Return the NetKet dual graph of the primal graph.

        The dual graph is typically constructed as the line graph of the primal graph and then converted to
        a :class:`netket.graph.Graph`.

        :return: NetKet dual graph.
        :raises NotImplementedError: If not implemented by the concrete graph class.
        """

        raise NotImplementedError

    @property
    @abstractmethod
    def nx_graph(self):
        r"""
        Return the NetworkX multigraph associated with the primal graph.

        :return: NetworkX primal graph.
        :raises NotImplementedError: If not implemented by the concrete graph class.
        """

        raise NotImplementedError

    @property
    @abstractmethod
    def dual_nx_graph(self):
        r"""
        Return the NetworkX graph associated with the dual of the primal graph.

        :return: NetworkX dual graph.
        :raises NotImplementedError: If not implemented by the concrete graph class.
        """

        raise NotImplementedError

    @property
    def signs(self):
        r"""
        Return local sign data for the primal graph.

        For a non planar embedding, signs are defined for ordered triples of edges incident at a vertex.
        Let :math:`t_1,t_2,t_3\in\mathbb R^3` be unit tangents of three incident oriented edges at a vertex
        :math:`v`. The sign is

        .. math::

           \varepsilon(e_1,e_2,e_3) = \operatorname{sgn}\bigl(\det[t_1\ t_2\ t_3]\bigr)

        For a planar embedding, signs reduce to ordered edge pairs using a 2D orientation convention. If
        :math:`t_1,t_2\in\mathbb R^2` are outgoing tangents at the vertex, an equivalent geometric quantity
        is the oriented area

        .. math::

           \operatorname{sgn}(t_{1x} t_{2y} - t_{1y} t_{2x})

        The handler decides which convention is used based on whether the graph is planar.

        :return: A dictionary of sign tables keyed by vertex.
        """

        return self.handler.graph_signs

    @property
    def is_planar(self):
        r"""
        Return True if the graph is planar under neuraLQX conventions.

        A planar graph uses integer vertex labels. A non planar graph uses 3 tuple vertex labels
        :math:`(x,y,z)` describing an embedding in :math:`\mathbb R^3`.

        :return: True if vertices are integer labelled, otherwise False.
        """

        return self.handler.is_planar

    @property
    def edges(self):
        r"""
        Return the list of oriented edges forming the graph.

        Each edge is represented as ``(u, v, key)`` where ``key`` distinguishes parallel edges.

        :return: List of oriented edges.
        """

        return self.handler.edges

    @property
    def vertices(self):
        r"""
        Return the list of vertices of the graph in the handler's internal representation.

        For planar graphs, vertices are integers. For non planar graphs, vertices are integers in the
        internal representation and the original coordinate tuples are accessible via the non planar
        mappings.

        :return: List of vertices.
        """

        return self.handler.vertices

    @property
    def nonplanar_edges(self):
        r"""
        Return the oriented edges in their original non planar coordinate representation.

        In the non planar case, input vertices may be 3 tuples :math:`(x,y,z)`. The handler stores a copy of
        the original edge list using these coordinate tuples. This property exposes that representation.

        :return: List of oriented edges with coordinate tuple endpoints.
        :raises NonExistentNonPlanarEdgesError: If called on a planar graph.
        """

        if self.is_planar:
            raise NonExistentNonPlanarEdgesError

        return self.handler._original_edges

    @property
    def mapping(self):
        r"""
        Return the bijective mapping from oriented edges to integer indices.

        Each oriented edge ``(u, v, key)`` is mapped to an integer index :math:`k`. This index identifies
        the position of the edge degree of freedom inside array based basis states, for example a local
        configuration :math:`\sigma\in\mathbb Z^N` where :math:`\sigma_k` stores the quantum number on the
        edge mapped to :math:`k`.

        :return: Dictionary mapping edges to indices.
        """

        return self.handler.mapping

    @property
    def nonplanar_vertex_mapping(self):
        r"""
        Return the mapping from non planar vertex coordinates to internal integer labels.

        For non planar graphs, the user may supply vertices as coordinate tuples :math:`(x,y,z)`. The
        handler relabels these to integers for efficient indexing. This property returns the forward map
        from coordinate tuples to integers.

        :return: Dictionary mapping coordinate tuples to integer labels.
        :raises NonExistentNonPlanarVerticesError: If called on a planar graph.
        """

        if self.is_planar:
            raise NonExistentNonPlanarVerticesError

        return {v: k for k, v in self.handler._nonplanar_vertex_map.items()}

    @property
    def n_edges(self):
        r"""
        Return the number of edges in the primal graph.

        :return: Number of primal edges.
        """

        return self.handler.nk_graph.n_edges

    @property
    def n_vertices(self):
        r"""
        Return the number of vertices in the primal graph.

        :return: Number of primal vertices.
        """

        return self.handler.nk_graph.n_nodes

    @property
    def n_dual_edges(self):
        r"""
        Return the number of edges in the dual graph.

        :return: Number of dual edges.
        """

        return self.handler.dual_nk_graph.n_edges

    @property
    def n_dual_vertices(self):
        r"""
        Return the number of vertices in the dual graph.

        :return: Number of dual vertices.
        """

        return self.handler.dual_nk_graph.n_nodes

    def edge_to_index(
        self,
        edge: Union[list, tuple],
    ) -> int:
        r"""
        Map an oriented edge to its integer index.

        The mapping is defined by the handler as a bijection between edges and indices. If the exact edge
        is not found, this method also tries the reversed orientation while keeping the same multiedge key.
        This supports situations where the caller provides an edge with opposite orientation to the stored
        convention.

        :param edge: Edge representation as a list or tuple. Expected form is ``(u, v, key)``.
        :return: Integer index of the edge degree of freedom.
        :raises InvalidIndexError: If neither the given edge nor its reversed orientation exists in the mapping.
        """

        # ensure tuple
        edge = tuple(edge)

        # note that when searching for the reversed edge, we reverse everything except the key
        idx = self.mapping.get(
            edge, self.mapping.get((*edge[:-1][::-1], edge[-1]), None)
        )

        if idx is None:
            raise InvalidIndexError(edge)

        return idx

    def index_to_edge(self, index: int) -> Union[list, tuple]:
        r"""
        Map an integer index back to an oriented edge.

        This is the inverse of :meth:`edge_to_index`. The index refers to the position of an edge degree of
        freedom in an array based configuration.

        :param index: Integer index to invert.
        :return: Edge triple ``(u, v, key)`` in the handler's orientation convention.
        :raises OutOfRangeIndexError: If the index is not present in the mapping.
        """

        # we are not doing multiple look-ups at the same time, so we will use this "brute" method
        # instead of reversing the dict and then using .get as in edge_to_index
        try:
            return list(
                list(self.mapping.keys())[list(self.mapping.values()).index(index)]
            )
        except ValueError:
            raise OutOfRangeIndexError(index, max(list(self.mapping.values())))

    def valence(self, node: Union[int, tuple]) -> int:
        r"""
        Return the valence of a vertex.

        The valence is the number of edges incident to the vertex in the underlying multigraph. In an
        oriented setting this counts both incoming and outgoing edges, with multiplicity for parallel
        edges.

        :param node: Vertex label, either an integer (planar) or a coordinate tuple (non planar input).
        :return: Valence of the vertex.
        """

        return len(list(self.handler.nx_graph.edges(node)))

    def nonplanar_vertex_to_index(self, nonplanar_vertex: tuple) -> int:
        r"""
        Convert a non planar coordinate vertex to its internal integer label.

        :param nonplanar_vertex: Vertex coordinate tuple :math:`(x,y,z)`.
        :return: Integer label used internally by the handler.
        :raises NonExistentNonPlanarVerticesError: If called on a planar graph.
        """

        if self.is_planar:
            raise NonExistentNonPlanarVerticesError

        return self.handler.convert_vertex_to_int(nonplanar_vertex)

    def export(
        self,
        filename: str,
        *,
        dpi: int = 300,
    ) -> None:
        r"""
        Export a visual representation of the graph to an image file.

        This delegates to the handler, which chooses an appropriate renderer depending on whether the graph
        is planar or non planar.

        :param filename: Output path for the exported image.
        :param dpi: Resolution in dots per inch used when rendering.
        :return: None.
        """

        self.handler.export(filename, dpi=dpi)

    def edges_at(self, vertex: int) -> list:
        r"""
        Return oriented edges whose head equals the specified vertex.

        This selects edges ``(u, v, key)`` such that ``v == vertex``. For a directed graph, these are the
        incoming edges at that vertex.

        :param vertex: Integer label of the vertex.
        :return: List of incoming oriented edges at `vertex`.
        """

        return [e for e in self.edges if e[1] == vertex]

    def edges_from(self, vertex: int) -> list:
        r"""
        Return oriented edges whose tail equals the specified vertex.

        This selects edges ``(u, v, key)`` such that ``u == vertex``. For a directed graph, these are the
        outgoing edges from that vertex.

        :param vertex: Integer label of the vertex.
        :return: List of outgoing oriented edges from `vertex`.
        """

        return [e for e in self.edges if e[0] == vertex]

    def __repr__(self):
        self.handler.__repr__()
        return ""

    def __str__(self):
        return self.handler.__str__()
