#!/usr/bin/env python3

from argparse import ArgumentParser
from pathlib import Path
from typing import TextIO

from ase import io, Atoms


# def assign_connectivity(atoms: Atoms) -> tuple[list[tuple[int, int, int]], list[int]]:
#     """Return bonds as (i, j, order) with 0-based atom indices and per-atom formal charges.

#     Bond orders follow V2000 bond types: 1 single, 2 double, 3 triple, 4 aromatic.
#     Dummy placeholder: no bonds, all atoms neutral.
#     """
#     return [], [0] * len(atoms)


def distance_connectivity(
    atoms: Atoms, tolerance: float = 1.2
) -> tuple[list[tuple[int, int, int]], list[int]]:
    import numpy as np
    from ase.data import covalent_radii

    """Distance baseline: bond when closer than tolerance * (sum of covalent radii), all single."""
    radii = covalent_radii[atoms.numbers]
    cutoff = tolerance * (radii[:, None] + radii[None, :])
    i, j = np.nonzero(np.triu(atoms.get_all_distances() < cutoff, k=1))
    return [(int(a), int(b), 1) for a, b in zip(i, j)], [0] * len(atoms)


def assign_connectivity(
    atoms: Atoms,
) -> tuple[list[tuple[int, int, int]], list[int]]:
    """Same contract as assign_connectivity, using RDKit's xyz2mol (rdDetermineBonds).

    Testing only: bonds are Kekulized; falls back to the distance baseline if RDKit fails.
    """
    from rdkit import Chem
    from rdkit.Chem import rdDetermineBonds

    xyz = f"{len(atoms)}\n\n" + "".join(
        f"{s} {x:.6f} {y:.6f} {z:.6f}\n"
        for s, (x, y, z) in zip(atoms.get_chemical_symbols(), atoms.positions)
    )
    try:
        mol = Chem.MolFromXYZBlock(xyz)
        rdDetermineBonds.DetermineBonds(
            mol, charge=int(atoms.info.get("charge", 0)), embedChiral=False
        )
        Chem.Kekulize(mol, clearAromaticFlags=True)
    except Exception:
        return distance_connectivity(atoms)

    bonds = [
        (b.GetBeginAtomIdx(), b.GetEndAtomIdx(), int(b.GetBondTypeAsDouble()))
        for b in mol.GetBonds()
    ]
    return bonds, [a.GetFormalCharge() for a in mol.GetAtoms()]


def write_sdf_record(
    fh: TextIO, atoms: Atoms, bonds: list[tuple[int, int, int]], charges: list[int]
) -> None:
    """Append one V2000 molfile record."""
    if len(atoms) > 999 or len(bonds) > 999:
        raise ValueError(f"{atoms.info.get('name')}: too large for V2000")

    fh.write(f"{atoms.info.get('name', '')}\n     assign         3D\n\n")
    fh.write(f"{len(atoms):3d}{len(bonds):3d}  0  0  0  0  0  0  0  0999 V2000\n")
    for sym, (x, y, z) in zip(atoms.get_chemical_symbols(), atoms.positions):
        fh.write(f"{x:10.4f}{y:10.4f}{z:10.4f} {sym:<3} 0" + "  0" * 11 + "\n")
    for i, j, order in bonds:
        fh.write(f"{i + 1:3d}{j + 1:3d}{order:3d}  0\n")

    charged = [(i + 1, q) for i, q in enumerate(charges) if q]
    for k in range(0, len(charged), 8):
        chunk = charged[k : k + 8]
        fh.write(
            f"M  CHG{len(chunk):3d}"
            + "".join(f" {i:3d} {q:3d}" for i, q in chunk)
            + "\n"
        )
    fh.write("M  END\n")
    fh.write("$$$$\n")


def main(input: Path, output: Path):
    with open(output, "w") as fh:
        for atoms in io.iread(input, index=slice(None)):
            bonds, charges = assign_connectivity(atoms)
            write_sdf_record(fh, atoms, bonds, charges)


if __name__ == "__main__":
    parser = ArgumentParser(
        description="Assign connectivity to xyz frames and write SDF."
    )
    parser.add_argument(
        "-i", "--input", type=Path, required=True, help="multi-frame (ext)xyz file"
    )
    parser.add_argument(
        "-o", "--output", type=Path, required=True, help="output SDF file"
    )
    args = parser.parse_args()
    main(args.input, args.output)