#!/usr/bin/env python3
"""Convert ring-fragments.txt and rigid-fragments.txt into compact C headers.

Reads:
  data/ring-fragments.txt   -- SMARTS pattern + xyz lines per fragment
  data/rigid-fragments.txt  -- SMILES + (atomic_num, xyz) lines per fragment

Writes:
  data/ringfragments.h
  data/rigidfragments.h

Each header exposes a flat `float kCoords[]` (xyz triplets concatenated) and
a `Fragment[]` table holding {pattern, coord_offset, natoms}. The runtime
loader can iterate the table without touching the filesystem.

The atomic-number column in rigid-fragments.txt is discarded -- the existing
builder.cpp loader reads it but never uses it (the SMILES already constrains
each atom's element).
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Iterable, TextIO


def parse_ring_fragments(path: Path) -> list[tuple[str, list[tuple[float, float, float]]]]:
    """Return [(smarts, [(x,y,z), ...]), ...] from ring-fragments.txt."""
    fragments: list[tuple[str, list[tuple[float, float, float]]]] = []
    smarts: str | None = None
    coords: list[tuple[float, float, float]] = []
    with path.open() as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            tokens = line.split()
            if len(tokens) == 1:
                if smarts is not None:
                    fragments.append((smarts, coords))
                smarts = tokens[0]
                coords = []
            elif len(tokens) == 3:
                coords.append((float(tokens[0]), float(tokens[1]), float(tokens[2])))
            else:
                raise ValueError(f"unexpected line in {path}: {line!r}")
    if smarts is not None:
        fragments.append((smarts, coords))
    return fragments


def parse_rigid_fragments(path: Path) -> list[tuple[str, list[tuple[float, float, float]]]]:
    """Return [(smiles, [(x,y,z), ...]), ...] from rigid-fragments.txt."""
    fragments: list[tuple[str, list[tuple[float, float, float]]]] = []
    smiles: str | None = None
    coords: list[tuple[float, float, float]] = []
    with path.open() as f:
        for line in f:
            stripped = line.strip()
            if not stripped:
                continue
            tokens = stripped.split()
            if len(tokens) == 1:
                if smiles is not None:
                    fragments.append((smiles, coords))
                smiles = tokens[0]
                coords = []
            elif len(tokens) == 4:
                # atomic_num, x, y, z -- discard atomic_num; the SMILES has it.
                coords.append((float(tokens[1]), float(tokens[2]), float(tokens[3])))
            else:
                raise ValueError(f"unexpected line in {path}: {stripped!r}")
    if smiles is not None:
        fragments.append((smiles, coords))
    return fragments


def cstring(s: str) -> str:
    # SMILES uses backslashes for cis/trans stereo (e.g. C=C\2); escape them.
    # No other special chars appear in the source data.
    return '"' + s.replace("\\", "\\\\").replace('"', '\\"') + '"'


def write_preamble(
    out: TextIO,
    guard: str,
    source_file: str,
    nfrag: int,
    natoms: int,
    includes: list[str],
) -> None:
    out.write(
        "/* Auto-generated by data/genfragmentheaders.py from {src}.\n"
        " * Do not edit by hand -- re-run the generator if {src} changes.\n"
        " *\n"
        " * {nfrag} fragments, {natoms} atoms total.\n"
        " * Coordinates stored as float (3 dp source precision).\n"
        " */\n".format(src=source_file, nfrag=nfrag, natoms=natoms)
    )
    out.write(f"#ifndef OB_{guard}_H\n")
    out.write(f"#define OB_{guard}_H\n\n")
    for inc in includes:
        out.write(f"#include <{inc}>\n")
    out.write("\nnamespace OpenBabel {\n\n")


def write_coords(out: TextIO, coords: list[tuple[float, float, float]]) -> None:
    out.write("{")
    for i, (x, y, z) in enumerate(coords):
        sep = ",\n      " if i > 0 else ""
        out.write(f"{sep}{x:.4f}f, {y:.4f}f, {z:.4f}f")
    out.write("}")


def emit_vector_header(
    out: TextIO,
    guard: str,
    table_name: str,
    struct_name: str,
    pattern_field: str,
    fragments: list[tuple[str, list[tuple[float, float, float]]]],
    source_file: str,
) -> None:
    """Emit fragments as a vector<Fragment>, preserving the source-file order.
    Use this when iteration order is semantically meaningful (e.g. ring
    fragments are pre-sorted by SMARTS complexity for match priority)."""
    total_atoms = sum(len(c) for _, c in fragments)
    write_preamble(out, guard, source_file, len(fragments), total_atoms,
                   ["string", "vector"])
    out.write(f"struct {struct_name} {{\n")
    out.write(f"  std::string {pattern_field};\n")
    out.write("  std::vector<float> coords;  // x,y,z triplets\n")
    out.write("};\n\n")

    out.write(f"inline const std::vector<{struct_name}> &{table_name}()\n")
    out.write("{\n")
    out.write(f"  static const std::vector<{struct_name}> table = {{\n")
    for pattern, coords in fragments:
        out.write(f"    {{{cstring(pattern)},\n")
        out.write("     ")
        write_coords(out, coords)
        out.write("},\n")
    out.write("  };\n  return table;\n}\n\n")
    out.write("} // namespace OpenBabel\n\n")
    out.write(f"#endif // OB_{guard}_H\n")


def emit_map_header(
    out: TextIO,
    guard: str,
    table_name: str,
    fragments: list[tuple[str, list[tuple[float, float, float]]]],
    source_file: str,
) -> None:
    """Emit fragments as a map<string, vector<float>>, keyed by pattern.
    Use this when the access pattern is exact-key lookup and iteration order
    doesn't matter (e.g. rigid fragments are looked up by exact SMILES)."""
    total_atoms = sum(len(c) for _, c in fragments)
    write_preamble(out, guard, source_file, len(fragments), total_atoms,
                   ["map", "string", "vector"])
    out.write(
        "inline const std::map<std::string, std::vector<float>> &"
        f"{table_name}()\n"
    )
    out.write("{\n")
    out.write(
        "  static const std::map<std::string, std::vector<float>> table = {\n"
    )
    for pattern, coords in fragments:
        out.write(f"    {{{cstring(pattern)},\n")
        out.write("     ")
        write_coords(out, coords)
        out.write("},\n")
    out.write("  };\n  return table;\n}\n\n")
    out.write("} // namespace OpenBabel\n\n")
    out.write(f"#endif // OB_{guard}_H\n")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--data-dir",
        type=Path,
        default=Path(__file__).resolve().parent,
        help="Directory containing the *-fragments.txt source files (default: same as this script).",
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=None,
        help="Output directory for the headers (default: same as --data-dir).",
    )
    args = parser.parse_args(argv)
    out_dir = args.out_dir or args.data_dir

    ring_src = args.data_dir / "ring-fragments.txt"
    rigid_src = args.data_dir / "rigid-fragments.txt"
    ring_dst = out_dir / "ringfragments.h"
    rigid_dst = out_dir / "rigidfragments.h"

    print(f"Parsing {ring_src}...", file=sys.stderr)
    ring = parse_ring_fragments(ring_src)
    print(f"  {len(ring)} ring fragments, {sum(len(c) for _, c in ring)} atoms", file=sys.stderr)

    print(f"Parsing {rigid_src}...", file=sys.stderr)
    rigid = parse_rigid_fragments(rigid_src)
    print(f"  {len(rigid)} rigid fragments, {sum(len(c) for _, c in rigid)} atoms", file=sys.stderr)

    print(f"Writing {ring_dst}...", file=sys.stderr)
    with ring_dst.open("w") as f:
        emit_vector_header(
            f,
            guard="RINGFRAGMENTS",
            table_name="RingFragments",
            struct_name="RingFragment",
            pattern_field="smarts",
            fragments=ring,
            source_file="ring-fragments.txt",
        )

    print(f"Writing {rigid_dst}...", file=sys.stderr)
    with rigid_dst.open("w") as f:
        emit_map_header(
            f,
            guard="RIGIDFRAGMENTS",
            table_name="RigidFragments",
            fragments=rigid,
            source_file="rigid-fragments.txt",
        )

    return 0


if __name__ == "__main__":
    sys.exit(main())
