"""ECFP-style circular fingerprints for COGILES colored graphs.

Generates fixed-length bit vector representations by iteratively hashing
each node's color together with its neighborhood, then folding all
collected identifiers into a bit vector via modulo.

Usage::

    import cogiles

    G = cogiles.parse("R-1RRRRR-1")

    # Quick usage with sensible defaults (radius=2, nbits=2048)
    fp = cogiles.fingerprint(G)

    # Configurable generator (reusable across many graphs)
    gen = cogiles.FingerprintGenerator(radius=3, nbits=1024)
    fp1 = gen.fingerprint(G)
    fp2 = gen.fingerprint(other_graph)

    # Jaccard/Tanimoto similarity (accepts strings, graphs, or fingerprints)
    sim = cogiles.jaccard_similarity("R-1RRRRR-1", G)

"""

from __future__ import annotations

import hashlib
from typing import Union

import networkx as nx

from cogiles.colors import rgb_to_symbol, SYMBOL_MAP


def _hash_to_int(*parts: str) -> int:
    """Deterministic hash of string parts to a non-negative integer.

    Uses BLAKE2b for determinism across Python sessions (unlike built-in
    ``hash()`` which is randomized by default).
    """
    h = hashlib.blake2b(digest_size=8)
    for part in parts:
        h.update(part.encode("utf-8"))
    return int.from_bytes(h.digest(), "big")


def _node_symbol(graph: nx.Graph, node: int) -> str:
    """Extract the canonical color symbol for a node.

    Mirrors the encoder's lookup order: use ``symbol`` if it's a
    recognized entry in the color registry, otherwise fall back to
    ``rgb_to_symbol`` on the ``color`` attribute.
    """
    symbol = graph.nodes[node].get("symbol")
    if symbol is not None and symbol in SYMBOL_MAP:
        return SYMBOL_MAP[symbol].symbol
    color = graph.nodes[node].get("color")
    if color is None:
        raise ValueError(
            f"Node {node} has no 'color' or valid 'symbol' attribute"
        )
    return rgb_to_symbol(tuple(color))


class FingerprintGenerator:
    """Configurable generator for ECFP-style colored-graph fingerprints.

    Parameters
    ----------
    radius : int
        Number of neighborhood expansion iterations. Radius 2 captures
        substructures up to 2 bonds from each center node (equivalent to
        ECFP4 diameter). Default: 2.
    nbits : int
        Length of the output bit vector. Larger values reduce hash
        collisions at the cost of sparser vectors. Default: 2048.
    """

    def __init__(self, *, radius: int = 2, nbits: int = 2048):
        if radius < 0:
            raise ValueError(f"radius must be >= 0, got {radius}")
        if nbits < 1:
            raise ValueError(f"nbits must be >= 1, got {nbits}")
        self.radius = radius
        self.nbits = nbits

    def fingerprint(self, graph: nx.Graph) -> list[int]:
        """Generate a bit-vector fingerprint for the given graph.

        Parameters
        ----------
        graph : nx.Graph
            A COGILES-style graph where each node has a ``symbol``
            and/or ``color`` attribute.

        Returns
        -------
        list[int]
            A list of ``nbits`` integers (0 or 1).
        """
        if graph.number_of_nodes() == 0:
            return [0] * self.nbits

        # Step 1: Initialize — each node gets an identifier from its symbol.
        identifiers: dict = {
            n: _hash_to_int(_node_symbol(graph, n)) for n in graph.nodes()
        }

        # Collect all identifiers across all radii (radius 0 = initial).
        all_identifiers: list[int] = list(identifiers.values())

        # Step 2: Iterative neighborhood expansion.
        for _ in range(self.radius):
            new_identifiers: dict = {}
            for node in graph.nodes():
                neighbor_ids = sorted(
                    str(identifiers[nb]) for nb in graph.neighbors(node)
                )
                new_identifiers[node] = _hash_to_int(
                    str(identifiers[node]), *neighbor_ids
                )
            identifiers = new_identifiers
            all_identifiers.extend(identifiers.values())

        # Step 3: Fold into bit vector.
        bits = [0] * self.nbits
        for ident in all_identifiers:
            bits[ident % self.nbits] = 1

        return bits

    def __repr__(self) -> str:
        return (
            f"FingerprintGenerator(radius={self.radius}, nbits={self.nbits})"
        )


# Module-level default generator (lazily shared).
_DEFAULT_GENERATOR = FingerprintGenerator(radius=2, nbits=2048)


def fingerprint(
    graph: nx.Graph, *, radius: int = 2, nbits: int = 2048
) -> list[int]:
    """Generate an ECFP-style fingerprint for a COGILES colored graph.

    Convenience function that creates or reuses a
    :class:`FingerprintGenerator` with the given settings.

    Parameters
    ----------
    graph : nx.Graph
        A COGILES-style graph where each node has a ``symbol``
        and/or ``color`` attribute.
    radius : int
        Neighborhood expansion depth. Default: 2.
    nbits : int
        Bit vector length. Default: 2048.

    Returns
    -------
    list[int]
        A list of ``nbits`` integers (0 or 1).
    """
    global _DEFAULT_GENERATOR
    if radius == _DEFAULT_GENERATOR.radius and nbits == _DEFAULT_GENERATOR.nbits:
        return _DEFAULT_GENERATOR.fingerprint(graph)

    gen = FingerprintGenerator(radius=radius, nbits=nbits)
    return gen.fingerprint(graph)


# Type alias for the accepted input types.
GraphLike = Union[str, nx.Graph, list[int]]


def _to_fingerprint(
    value: GraphLike, *, radius: int, nbits: int
) -> list[int]:
    """Coerce a COGILES string, nx.Graph, or pre-computed fingerprint.

    Defers the ``cogiles.parse`` import to avoid circular imports.
    """
    if isinstance(value, list):
        return value
    if isinstance(value, str):
        from cogiles.parser import parse

        value = parse(value)
    if isinstance(value, nx.Graph):
        return fingerprint(value, radius=radius, nbits=nbits)
    raise TypeError(
        f"Expected a COGILES string, nx.Graph, or list[int], "
        f"got {type(value).__name__}"
    )


def jaccard_similarity(
    a: GraphLike,
    b: GraphLike,
    *,
    radius: int = 2,
    nbits: int = 2048,
) -> float:
    """Jaccard (Tanimoto) similarity between two colored graphs.

    Computes circular fingerprints for both inputs and returns their
    Jaccard similarity: ``|A & B| / |A | B|``.

    Parameters
    ----------
    a, b : str | nx.Graph | list[int]
        Each input can be a COGILES string, an ``nx.Graph``, or a
        pre-computed fingerprint (list of 0/1 ints).
    radius : int
        Neighborhood expansion depth for fingerprint generation.
        Ignored when a pre-computed fingerprint is passed. Default: 2.
    nbits : int
        Bit vector length for fingerprint generation.
        Ignored when a pre-computed fingerprint is passed. Default: 2048.

    Returns
    -------
    float
        Similarity score in [0.0, 1.0]. Returns 1.0 when both
        fingerprints are all-zero (both graphs are empty).
    """
    fp_a = _to_fingerprint(a, radius=radius, nbits=nbits)
    fp_b = _to_fingerprint(b, radius=radius, nbits=nbits)

    if len(fp_a) != len(fp_b):
        raise ValueError(
            f"Fingerprint lengths differ: {len(fp_a)} vs {len(fp_b)}"
        )

    intersection = sum(x & y for x, y in zip(fp_a, fp_b))
    union = sum(x | y for x, y in zip(fp_a, fp_b))

    if union == 0:
        return 1.0

    return intersection / union
