"""Reader for the Reaction-QM B3LYP-RXN subset.

Reaction-QM (KAIST InnoCORE AI-CRED Institute + Samsung Electronics
Semiconductor R&D Center, *Sci Data* 2026, doi:10.1038/s41597-026-07325-w;
data at doi:10.5281/zenodo.18551029, CC-BY-4.0) ships three subsets. This
module reads only **B3LYP-RXN**: 199,890 reactions optimized at
B3LYP-D3/TZVP, each with fully optimized reactant, product *and transition
state* geometries.

Every reaction carries a reference TS and barrier, so a TS found by
`mepd channels` can be scored against it.

On-disk layout of `B3LYPD3_TZVP.h5`, which the authors' `open_h5_example.py`
documents by example rather than by spec::

    <chunk group>/              e.g. "RXN_0000000001_to_RXN_0000010000"
        RXN_0000000001/
            R0/ R1/ ...         reactant fragments
            P0/ P1/ ...         product fragments
            TS/                 (its `smiles` is the full reaction SMILES)

and every species group holds `smiles`, `EHG` (Hartree), `charge`,
`multiplicity`, `atomic_numbers`, `coordinates` (Angstrom).

Reactions are keyed inside chunk groups, so locating one means knowing its
chunk. `build_index` walks the file once and caches that mapping next to the
h5; everything else goes through it rather than rescanning.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator, Optional

DATA_DIR = Path(__file__).resolve().parents[2] / "data" / "reaction_qm"
H5_PATH = DATA_DIR / "B3LYPD3_TZVP.h5"
INFO_CSV = DATA_DIR / "B3LYPD3_TZVP_reaction_info.csv"
INDEX_PATH = DATA_DIR / "B3LYPD3_TZVP_index.json"

HARTREE_TO_KCAL = 627.5094740631

_MAPNUM_RE = re.compile(r":(\d+)\]")

# Z -> symbol, covering exactly the elements Reaction-QM spans.
_SYMBOLS = {
    1: "H", 6: "C", 7: "N", 8: "O", 9: "F",
    14: "Si", 15: "P", 16: "S", 17: "Cl",
}


@dataclass
class Species:
    """One reactant fragment, product fragment, or transition state."""

    tag: str                  # "R0", "P1", "TS", ...
    smiles: str
    energy: float             # E, Hartree
    enthalpy: float           # H, Hartree
    gibbs: float              # G, Hartree
    charge: int
    multiplicity: int
    atomic_numbers: list
    coordinates: list         # [[x, y, z], ...] in Angstrom

    @property
    def symbols(self) -> list:
        return [_SYMBOLS.get(int(z), f"Z{int(z)}") for z in self.atomic_numbers]

    @property
    def n_atoms(self) -> int:
        return len(self.atomic_numbers)

    @property
    def n_heavy(self) -> int:
        return sum(1 for z in self.atomic_numbers if int(z) != 1)

    def to_xyz(self, comment: str = "") -> str:
        lines = [str(self.n_atoms), comment or self.smiles]
        for sym, (x, y, z) in zip(self.symbols, self.coordinates):
            lines.append(f"{sym:<3} {x:>18.12f} {y:>18.12f} {z:>18.12f}")
        return "\n".join(lines) + "\n"


@dataclass
class Reaction:
    """One B3LYP-RXN entry: reactant side, product side, and the TS joining
    them. Fragments stay separate -- which side is multi-fragment is exactly
    what makes a reaction bimolecular, and collapsing that would lose it."""

    reaction_id: str
    reactants: list           # list[Species]
    products: list            # list[Species]
    ts: Optional[Species]

    @property
    def reaction_smiles(self) -> str:
        """The mapped reaction SMILES, which the TS group carries."""
        return self.ts.smiles if self.ts is not None else ""

    @property
    def start_smiles(self) -> str:
        return ".".join(s.smiles for s in self.reactants)

    @property
    def end_smiles(self) -> str:
        return ".".join(s.smiles for s in self.products)

    @property
    def charge(self) -> int:
        return sum(s.charge for s in self.reactants)

    @property
    def n_atoms(self) -> int:
        return sum(s.n_atoms for s in self.reactants)

    @property
    def n_heavy(self) -> int:
        return sum(s.n_heavy for s in self.reactants)

    def barrier_kcal(self) -> Optional[float]:
        """Electronic barrier from the stored energies, in kcal/mol.

        Computed from the separated-fragment sum, which is the same
        convention the dataset's own `dE_dagger` column uses -- for a
        bimolecular reaction this is the barrier relative to infinitely
        separated reactants, not to a pre-reaction complex."""
        if self.ts is None or not self.reactants:
            return None
        e_r = sum(s.energy for s in self.reactants)
        return (self.ts.energy - e_r) * HARTREE_TO_KCAL

    def is_balanced(self) -> bool:
        """Same multiset of elements on both sides. Reaction-QM reactions are
        atom-mapped elementary steps and should all pass; checking is cheap
        and a failure would mean the entry is unusable as an endpoint pair."""
        from collections import Counter

        left = Counter(z for s in self.reactants for z in map(int, s.atomic_numbers))
        right = Counter(z for s in self.products for z in map(int, s.atomic_numbers))
        return left == right

    def write_endpoints(self, outdir: Path) -> dict:
        """Write reactant/product/TS geometries as xyz files.

        Feeding real B3LYP-optimized geometries rather than SMILES is the
        main reason to bother with the h5 at all: it avoids re-embedding a
        multi-fragment reactant side, which is where separated fragments can
        relax into a spuriously bonded adduct and make the requested
        reaction unfindable by construction.

        **Atom ordering is not shared between the two endpoints.** Each
        species group stores its own SMILES with map numbers restarting at 1
        per fragment, and its xyz follows that local order, so concatenating
        fragments gives sides that are ordered differently from one another
        (72 of the 100 benchmark reactions). The reaction's *global* map
        lives only on the TS group's reaction SMILES, and no rule tried so
        far reproduces local -> global on more than about half the set:
        neither positional alignment of map numbers between the local and
        global fragment strings (38/100), nor "local map k is the k-th
        smallest global map in that fragment" (52/100). Rather than ship a
        mapping that is silently wrong half the time, these files are
        written in natural per-fragment order and the endpoint
        correspondence is left to `mepd channels --atom-mapping`, which
        exists for exactly this case. The TS reference *is* in global map
        order, since the TS is a single structure.

        **Fragments are pulled apart before writing.** Each species group is
        stored in its own coordinate frame, centred near the origin, so
        concatenating a multi-fragment side verbatim stacks the fragments on
        top of one another -- a three-fragment product side came out as an
        oxygen with seven bonds the first time this was run. Fragments are
        therefore packed (random orientations in a small sphere, see
        `_separate_fragments`); for the benchmark, `dock_endpoints` docks them
        with xtb aISS instead. Their internal geometry is untouched; only the
        placement is synthetic, and the endpoint minimization will settle it.
        """
        outdir = Path(outdir)
        outdir.mkdir(parents=True, exist_ok=True)
        written = {}
        for name, species in (("start", self.reactants), ("end", self.products)):
            atoms = _separate_fragments(species)
            lines = [str(len(atoms)), f"{self.reaction_id} {name}: " + ".".join(s.smiles for s in species)]
            for sym, (x, y, z) in atoms:
                lines.append(f"{sym:<3} {x:>18.12f} {y:>18.12f} {z:>18.12f}")
            fp = outdir / f"{name}.xyz"
            fp.write_text("\n".join(lines) + "\n")
            written[name] = fp
        if self.ts is not None:
            fp = outdir / "ts_reference.xyz"
            fp.write_text(self.ts.to_xyz(f"{self.reaction_id} reference TS (B3LYP-D3/TZVP)"))
            written["ts"] = fp
        return written


def _separate_fragments(species: list, seed: int = 0) -> list:
    """The fragments of one side packed together (mepd.complexes.packed:
    random orientations in a small sphere, 2 Angstrom apart at their
    closest), returning [(symbol, (x, y, z)), ...] in fragment order. Never a
    line-up: fragments in a row give the path search a meaningless start.

    A single fragment is returned untouched -- its stored geometry is the
    real optimized one."""
    if len(species) == 1:
        sp = species[0]
        return list(zip(sp.symbols, [tuple(c) for c in sp.coordinates]))
    import numpy as np
    from qcdata import Structure

    from mepd.complexes import packed

    parts = [Structure(symbols=list(sp.symbols), geometry=np.asarray(sp.coordinates) * 1.8897259886,
                       charge=0, multiplicity=1) for sp in species]
    s = packed(parts, seed=seed)
    return list(zip(s.symbols, [tuple(c) for c in np.asarray(s.geometry).reshape(-1, 3) / 1.8897259886]))


def _decode(value) -> str:
    return value.decode() if isinstance(value, bytes) else str(value)


def _read_species(tag: str, group) -> Species:
    ehg = list(group["EHG"][()])
    return Species(
        tag=tag,
        smiles=_decode(group["smiles"][()]),
        energy=float(ehg[0]),
        enthalpy=float(ehg[1]),
        gibbs=float(ehg[2]),
        charge=int(group["charge"][()]),
        multiplicity=int(group["multiplicity"][()]),
        atomic_numbers=[int(z) for z in group["atomic_numbers"][()]],
        coordinates=[[float(c) for c in xyz] for xyz in group["coordinates"][()]],
    )


def build_index(h5_path: Path = H5_PATH, cache: Optional[Path] = INDEX_PATH) -> dict:
    """Map every reaction id to the chunk group holding it, caching the
    result. Without this, looking up a reaction means scanning chunk groups
    until one of them contains the key."""
    import h5py

    cache = Path(cache) if cache else None
    if cache and cache.is_file():
        return json.loads(cache.read_text())

    index = {}
    with h5py.File(h5_path, "r") as fh:
        for chunk_name, chunk in fh.items():
            for rxn_id in chunk.keys():
                index[rxn_id] = chunk_name
    if cache:
        cache.write_text(json.dumps(index))
    return index


def load_reaction(handle, rxn_id: str, index: dict) -> Reaction:
    """Read one reaction from an open `h5py.File`."""
    if rxn_id.isdigit():
        rxn_id = f"RXN_{rxn_id.rjust(10, '0')}"
    chunk = index.get(rxn_id)
    if chunk is None:
        raise KeyError(f"{rxn_id} not present in the B3LYP-RXN index")

    group = handle[chunk][rxn_id]
    reactants, products, ts = [], [], None
    for tag in sorted(group.keys()):
        species = _read_species(tag, group[tag])
        if tag == "TS":
            ts = species
        elif tag.startswith("R"):
            reactants.append(species)
        elif tag.startswith("P"):
            products.append(species)
    return Reaction(reaction_id=rxn_id, reactants=reactants, products=products, ts=ts)


def iter_reactions(
    h5_path: Path = H5_PATH, limit: Optional[int] = None
) -> Iterator[Reaction]:
    """Stream reactions in file order -- the whole subset is far too large to
    hold in memory as parsed objects."""
    import h5py

    index = build_index(h5_path)
    n = 0
    with h5py.File(h5_path, "r") as fh:
        for chunk_name, chunk in fh.items():
            for rxn_id in chunk.keys():
                yield load_reaction(fh, rxn_id, index)
                n += 1
                if limit is not None and n >= limit:
                    return