import numpy as np
import hashlib

# ---------------------------
# Fast uint64 hashing helpers
# ---------------------------

_U64 = np.uint64

def _mix_u64(x: np.ndarray) -> np.ndarray:
    """
    SplitMix64 mixer (vectorized). Input/Output are uint64 arrays.
    Produces well-distributed 64-bit values.
    """
    x = (x + _U64(0x9E3779B97F4A7C15)) & _U64(0xFFFFFFFFFFFFFFFF)
    x = (x ^ (x >> _U64(30))) * _U64(0xBF58476D1CE4E5B9) & _U64(0xFFFFFFFFFFFFFFFF)
    x = (x ^ (x >> _U64(27))) * _U64(0x94D049BB133111EB) & _U64(0xFFFFFFFFFFFFFFFF)
    x = x ^ (x >> _U64(31))
    return x

def _hash_pair_u64(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """
    Order-sensitive pair hash: h(a,b). Vectorized.
    """
    return _mix_u64(a ^ _mix_u64(b + _U64(0xD1B54A32D192ED03)))

def _sha256_hex_from_u64(arr_u64: np.ndarray) -> str:
    """
    Final stable cryptographic hash from uint64 array.
    """
    return hashlib.sha256(arr_u64.tobytes()).hexdigest()


# -----------------------------------------
# Periodic minimum-image (vectorized, fast)
# -----------------------------------------

def minimum_image_vectors_and_shifts(dr: np.ndarray, lattice: np.ndarray):
    """
    dr: (M,3) Cartesian displacement vectors (rj-ri) in same basis as lattice.
    lattice: (3,3) with columns = lattice vectors (consistent with your tree.bounds usage).

    Returns:
      dr_mic: (M,3) minimum-image displacement in Cartesian
      shift_i: (M,3) integer lattice shifts applied (so dr_mic = dr - shift_i @ lattice)
    """
    # Use solve with lattice.T because your code uses shifts_i.dot(bounds) with bounds = lattice
    # dr = frac @ lattice  <=>  frac = dr @ inv(lattice)
    inv_lat = np.linalg.inv(lattice)
    frac = dr @ inv_lat                     # (M,3)
    shift_i = np.rint(frac).astype(np.int32) # nearest integer shifts
    frac_mic = frac - shift_i
    dr_mic = frac_mic @ lattice
    return dr_mic, shift_i


# -----------------------------------------
# Efficient periodic chemical graph builder
# -----------------------------------------

def build_periodic_chemical_graph(
    tree,
    species,
    covalent_radii,
    lattice=None,
    tol=1.2,
    dist_bin=0.02,
    include_shift=True,
    include_dist=True,
):
    """
    Build a periodic neighbor graph using:
      cutoff_ij = tol * (R_i + R_j),
    with minimum-image vectors and (optional) lattice shift integers.

    Parameters
    ----------
    tree : PeriodicCKDTree
        Your KDTree object (already PBC-aware).
    species : (N,) array-like
        Species labels for each atom. Can be ints (atomic numbers) or strings ("Ni", "O", ...).
    covalent_radii : dict
        Map species label -> covalent radius (Å). Keys must match entries in `species`.
        Example: {"Ni":1.24, "O":0.66, "H":0.31, "K":2.03}
        or {28:1.24, 8:0.66, 1:0.31, 19:2.03}
    lattice : (3,3) array-like, optional
        Lattice matrix. If None, uses tree.bounds.
        IMPORTANT: must be consistent with how your positions are expressed.
    tol : float
        Multiplicative factor for bond cutoff.
    dist_bin : float
        Distance bin width (Å). Edge label uses integer bin = round(d / dist_bin).
        This makes the graph robust to tiny relaxations.
    include_shift : bool
        If True, include periodic image integer shift (sx,sy,sz) in edges.
    include_dist : bool
        If True, include binned distance in edges.

    Returns
    -------
    graph : dict
        {
          "N": int,
          "species": np.ndarray,
          "edges_i": np.ndarray (E,),
          "edges_j": np.ndarray (E,),
          "dist_bin": np.ndarray (E,)  (optional),
          "shift": np.ndarray (E,3)    (optional),
        }
    """
    pos = np.asarray(tree.data, float)
    species = np.asarray(species)
    N = pos.shape[0]
    if lattice is None:
        lattice = np.asarray(tree.bounds, float)
    lattice = np.asarray(lattice, float)
    if lattice.shape != (3, 3):
        raise ValueError("lattice must be shape (3,3)")

    # Radii array aligned with atoms (vectorized lookup)
    # Fast path: species are hashable; build radii per atom once.
    try:
        r_atom = np.array([covalent_radii[s] for s in species], dtype=float)
    except KeyError as e:
        raise KeyError(f"Missing covalent radius for species: {e}")

    # Global maximum cutoff to get candidate pairs once
    r_max = tol * (2.0 * float(np.max(r_atom)))
    pairs = tree.query_pairs(r_max)  # list of (i,j)
    if len(pairs) == 0:
        graph = {"N": N, "species": species, "edges_i": np.array([], dtype=np.int32), "edges_j": np.array([], dtype=np.int32)}
        if include_dist:
            graph["dist_bin"] = np.array([], dtype=np.int16)
        if include_shift:
            graph["shift"] = np.zeros((0, 3), dtype=np.int16)
        return graph

    pairs = np.asarray(pairs, dtype=np.int32)  # (M,2)
    i = pairs[:, 0]
    j = pairs[:, 1]

    # Vectorized minimum-image displacement and shift integers
    dr = pos[j] - pos[i]
    dr_mic, shift_i = minimum_image_vectors_and_shifts(dr, lattice)
    d = np.linalg.norm(dr_mic, axis=1)

    # Chemical cutoff per pair (vectorized)
    cutoff_ij = tol * (r_atom[i] + r_atom[j])
    keep = d <= cutoff_ij
    if not np.any(keep):
        graph = {"N": N, "species": species, "edges_i": np.array([], dtype=np.int32), "edges_j": np.array([], dtype=np.int32)}
        if include_dist:
            graph["dist_bin"] = np.array([], dtype=np.int16)
        if include_shift:
            graph["shift"] = np.zeros((0, 3), dtype=np.int16)
        return graph

    i = i[keep]
    j = j[keep]
    d = d[keep]
    shift_i = shift_i[keep]

    # Edge distance label (binned, integer)
    if include_dist:
        db = np.rint(d / float(dist_bin)).astype(np.int16)

    # Store shift as small ints (typically -1/0/1 but can be larger)
    if include_shift:
        sh = shift_i.astype(np.int16)

    graph = {
        "N": N,
        "species": species,
        "edges_i": i,
        "edges_j": j,
    }
    if include_dist:
        graph["dist_bin"] = db
    if include_shift:
        graph["shift"] = sh

    return graph


# ---------------------------------------------------
# Canonical hashing via fast WL refinement (no NX)
# ---------------------------------------------------

def canonical_graph_hash_wl(
    graph,
    wl_iters=3,
    use_edge_dist=True,
    use_edge_shift=True,
    final_sha256=True,
):
    """
    Compute a canonical (permutation-invariant) hash of the graph using
    Weisfeiler–Lehman refinement with uint64 labels.

    This is fast and works well for chemical graphs. It is not a formal
    NAUTY-style proof-level canonical form, but in practice it is extremely
    discriminative for atomistic graphs.

    Parameters
    ----------
    graph : dict
        Output of build_periodic_chemical_graph.
    wl_iters : int
        Number of WL refinement iterations (2–4 usually enough).
    use_edge_dist : bool
        Include dist_bin in neighbor messages.
    use_edge_shift : bool
        Include periodic shift integers in neighbor messages.
    final_sha256 : bool
        If True, return SHA256(hex). If False, return uint64 signature array.

    Returns
    -------
    h : str or np.ndarray
        Canonical hash string (sha256) or uint64 signature array (sorted labels).
    """
    N = int(graph["N"])
    species = np.asarray(graph["species"])
    ei = np.asarray(graph["edges_i"], dtype=np.int32)
    ej = np.asarray(graph["edges_j"], dtype=np.int32)
    E = ei.size

    # Build undirected edge list in "directed" form for message passing:
    # src -> dst and dst -> src
    src = np.concatenate([ei, ej]).astype(np.int32)
    dst = np.concatenate([ej, ei]).astype(np.int32)
    M = src.size

    # Edge labels aligned with directed edges
    edge_label = None
    if use_edge_dist and "dist_bin" in graph:
        db = np.asarray(graph["dist_bin"], dtype=np.int16)
        db_dir = np.concatenate([db, db]).astype(np.int32)
    else:
        db_dir = None

    if use_edge_shift and "shift" in graph:
        sh = np.asarray(graph["shift"], dtype=np.int16)  # (E,3)
        # For reverse direction, shift should be negated
        sh_dir = np.vstack([sh, -sh]).astype(np.int32)    # (2E,3)
    else:
        sh_dir = None

    # Initial node labels from species (stable)
    # If species are strings, hash them deterministically to uint64
    if species.dtype.kind in ("U", "S", "O"):
        # Stable conversion: hash bytes of each species string via sha256 -> take first 8 bytes
        sbytes = np.array([hashlib.sha256(str(s).encode()).digest()[:8] for s in species], dtype="|S8")
        labels = np.frombuffer(sbytes.tobytes(), dtype=np.uint64)
    else:
        labels = species.astype(np.uint64)

    labels = _mix_u64(labels + _U64(0x123456789ABCDEF0))

    # Pre-sort edges by dst to aggregate neighbor messages quickly
    order = np.argsort(dst, kind="mergesort")
    src = src[order]
    dst = dst[order]
    if db_dir is not None:
        db_dir = db_dir[order]
    if sh_dir is not None:
        sh_dir = sh_dir[order]

    # Compute segment boundaries for each node in dst
    counts = np.bincount(dst, minlength=N)
    offsets = np.zeros(N + 1, dtype=np.int64)
    np.cumsum(counts, out=offsets[1:])

    # WL iterations
    for _ in range(int(wl_iters)):
        # Message for each directed edge: depends on neighbor label + edge attributes
        neigh_lab = labels[src]

        msg = neigh_lab
        if db_dir is not None:
            msg = _hash_pair_u64(msg, _mix_u64(db_dir.astype(np.uint64)))
        if sh_dir is not None:
            # pack shift (sx,sy,sz) into uint64 deterministically
            # shift values are small; bias + mask to keep non-negative
            sx = (sh_dir[:, 0].astype(np.int64) + 2048) & 0xFFF
            sy = (sh_dir[:, 1].astype(np.int64) + 2048) & 0xFFF
            sz = (sh_dir[:, 2].astype(np.int64) + 2048) & 0xFFF
            packed = (sx.astype(np.uint64) << _U64(24)) ^ (sy.astype(np.uint64) << _U64(12)) ^ sz.astype(np.uint64)
            msg = _hash_pair_u64(msg, _mix_u64(packed))

        msg = _mix_u64(msg)

        # Aggregate messages per destination node:
        # Use commutative aggregation (sum + xor) to avoid sorting neighbors (fast).
        new_labels = np.empty_like(labels)
        for v in range(N):
            a, b = offsets[v], offsets[v + 1]
            if a == b:
                agg = _U64(0)
                agg2 = _U64(0)
            else:
                mv = msg[a:b]
                agg = mv.sum(dtype=np.uint64)
                agg2 = np.bitwise_xor.reduce(mv)

            base = labels[v]
            # combine base label with aggregates
            combined = _hash_pair_u64(base, _mix_u64(agg ^ _mix_u64(agg2)))
            new_labels[v] = combined

        labels = _mix_u64(new_labels)

    # Canonical signature: sort final node labels (permutation-invariant)
    sig = np.sort(labels)

    if final_sha256:
        return _sha256_hex_from_u64(sig)
    return sig


# --------------------------
# One-call convenience API
# --------------------------

def canonical_periodic_graph_hash(
    tree,
    species,
    covalent_radii,
    lattice=None,
    tol=1.2,
    dist_bin=0.02,
    wl_iters=3,
    use_edge_shift=True,
    use_edge_dist=True,
):
    """
    Full pipeline:
      KDTree -> periodic chemical graph -> canonical WL hash (sha256)

    Returns:
      (hash_hex, graph_dict)
    """
    g = build_periodic_chemical_graph(
        tree=tree,
        species=species,
        covalent_radii=covalent_radii,
        lattice=lattice,
        tol=tol,
        dist_bin=dist_bin,
        include_shift=use_edge_shift,
        include_dist=use_edge_dist,
    )
    h = canonical_graph_hash_wl(
        g,
        wl_iters=wl_iters,
        use_edge_dist=use_edge_dist,
        use_edge_shift=use_edge_shift,
        final_sha256=True,
    )
    return h, g

import typing

def structure_graph_hash_factory(
    tol: float = 1.2,
    dist_bin: float = 0.02,
    wl_iters: int = 3,
    use_edge_shift: bool = True,
    use_edge_dist: bool = True,
) -> typing.Callable[[typing.Any], str]:
    """
    Returns a hash function that computes a canonical permutation-invariant 
    hash for a Structure object based on its periodic chemical graph.

    This is a factory function that binds the specified parameters and 
    returns a callable. The returned callable takes a Structure object,
    automatically extracts the necessary spatial and chemical information 
    (KDTree, species, covalent radii, lattice vectors) from its 
    AtomPositionManager, and processes it through a Weisfeiler-Lehman (WL) 
    graph hashing pipeline.

    Parameters
    ----------
    tol : float, optional
        Tolerance multiplier for the covalent bond cutoff distance. Default is 1.2.
    dist_bin : float, optional
        Bin width for discretizing edge distances (in Å). Default is 0.02.
    wl_iters : int, optional
        Number of WL refinement iterations to perform. Default is 3.
    use_edge_shift : bool, optional
        If True, includes periodic boundary shift information in edge messages. Default is True.
    use_edge_dist : bool, optional
        If True, includes the binned distance in edge messages. Default is True.

    Returns
    -------
    Callable[[Any], str]
        A function that takes a Structure object and returns a canonical 
        SHA-256 hash string representing its chemical graph.
    """
    def hash_fn(structure) -> str:
        # 1. Build the periodic chemical graph using the structure's internal properties
        g = build_periodic_chemical_graph(
            tree=structure.AtomPositionManager.kdtree,
            species=structure.AtomPositionManager.atomLabelsList,
            covalent_radii=structure.AtomPositionManager.covalent_radii,
            lattice=structure.AtomPositionManager.latticeVectors,
            tol=tol,
            dist_bin=dist_bin,
            include_shift=use_edge_shift,
            include_dist=use_edge_dist,
        )
        
        # 2. Compute canonical hash using Weisfeiler-Lehman (WL) refinement
        h = canonical_graph_hash_wl(
            g,
            wl_iters=wl_iters,
            use_edge_dist=use_edge_dist,
            use_edge_shift=use_edge_shift,
            final_sha256=True,
        )
        
        return h

    return hash_fn

# --------------------------
# Usage example (sketch)
# --------------------------
# radii = {"Ni": 1.24, "O": 0.66, "H": 0.31, "K": 2.03}
# h, graph = canonical_periodic_graph_hash(tree, species_array, radii, lattice=tree.bounds)
# print(h)
# graph["edges_i"], graph["edges_j"], graph.get("dist_bin"), graph.get("shift")


