"""Bond detection using CrystalNN."""

import numpy as np
from typing import Dict, List, Optional

from pymatgen.analysis.local_env import CrystalNN

from ceraviz.types import CrystalTD
from ceraviz.structure import build_pmg_structure


def find_bonds_crystalnn(
    structure_td: CrystalTD,
    oxidation_states: Optional[Dict[str, int]] = None,
) -> List[dict]:
    """Find bonds in a crystal structure using CrystalNN algorithm.

    Returns a list of bond dicts with keys: i, j, image, distance, weight.
    """
    pmg_structure = build_pmg_structure(structure_td, oxidation_states)

    cnn = CrystalNN()
    bonds = []
    seen = set()
    lattice = np.array(pmg_structure.lattice.matrix)

    for i in range(len(pmg_structure)):
        for nb in cnn.get_nn_info(pmg_structure, i):
            j = nb["site_index"]
            image = tuple(nb["image"])
            weight = float(nb.get("weight", 1.0))

            key = tuple(sorted((i, j))) + image
            if key in seen:
                continue
            seen.add(key)

            ri = pmg_structure[i].coords
            rj = pmg_structure[j].coords + np.array(image) @ lattice
            dist = float(np.linalg.norm(ri - rj))

            bonds.append({
                "i": i,
                "j": j,
                "image": image,
                "distance": dist,
                "weight": weight,
            })

    return bonds
