#!/usr/bin/env python3
"""Score assigned connectivity against the reference SDF by exact bond type.

Records are paired by position and must agree on name and atom count. Bonds are
keyed by unordered atom pair over the union of both bond sets; a pair present on
only one side gets type 0, so missing and extra bonds both count as mismatches.
Types are compared literally (aromatic 4 never matches a Kekule 1 or 2), except
that equivalent Kekule structures are not penalized: any ring whose bonds
alternate 1/2 in the output may be flipped (1 <-> 2), repeatedly and for rings
of any size or element, and the flipped assignment that best matches the
reference is the one scored. Scores are pooled over all records.
"""

from argparse import ArgumentParser
from collections import defaultdict
from itertools import zip_longest
from pathlib import Path

HERE = Path(__file__).parent


def read_records(path):
    with open(path) as fh:
        record = []
        for line in fh:
            if line.startswith("$$$$"):
                yield record
                record = []
            else:
                record.append(line.rstrip("\n"))
        if any(ln.strip() for ln in record):
            yield record


def parse_bonds(lines):
    """Return (name, n_atoms, {frozenset((i, j)): type}) for one V2000 record."""
    name = lines[0].strip()
    counts = lines[3]
    if "V3000" in counts:
        raise ValueError(f"{name}: V3000 not supported")
    n_atoms, n_bonds = int(counts[0:3]), int(counts[3:6])
    bonds = {}
    for line in lines[4 + n_atoms : 4 + n_atoms + n_bonds]:
        i, j, order = int(line[0:3]), int(line[3:6]), int(line[6:9])
        bonds[frozenset((i, j))] = order
    return name, n_atoms, bonds


def ring_bonds(bonds):
    """Return the single/double bonds that lie on a ring of single/double bonds.

    These are the non-bridge edges of the 1/2 bond graph (Tarjan, iterative).
    """
    adj = defaultdict(list)
    for pair, order in bonds.items():
        if order in (1, 2):
            i, j = pair
            adj[i].append((j, pair))
            adj[j].append((i, pair))
    disc, low, cyclic = {}, {}, set()
    for root in adj:
        if root in disc:
            continue
        disc[root] = low[root] = len(disc)
        stack = [(root, None, iter(adj[root]))]
        while stack:
            a, via, edges = stack[-1]
            for b, pair in edges:
                if pair == via:
                    continue
                if b in disc:
                    # a non-tree edge always closes a ring
                    cyclic.add(pair)
                    low[a] = min(low[a], disc[b])
                else:
                    disc[b] = low[b] = len(disc)
                    stack.append((b, pair, iter(adj[b])))
                    break
            else:
                stack.pop()
                if stack:
                    parent = stack[-1][0]
                    low[parent] = min(low[parent], low[a])
                    if low[a] <= disc[parent]:
                        cyclic.add(via)
    return cyclic


def ring_systems(pairs):
    """Split bond pairs into connected components, each listed in BFS order."""
    adj = defaultdict(list)
    for pair in pairs:
        i, j = pair
        adj[i].append(pair)
        adj[j].append(pair)
    seen_atoms, seen_pairs, systems = set(), set(), []
    for root in sorted(adj):
        if root in seen_atoms:
            continue
        seen_atoms.add(root)
        queue, system = [root], []
        for a in queue:
            for pair in adj[a]:
                if pair not in seen_pairs:
                    seen_pairs.add(pair)
                    system.append(pair)
                    (b,) = pair - {a}
                    if b not in seen_atoms:
                        seen_atoms.add(b)
                        queue.append(b)
        systems.append(system)
    return systems


def best_system(pairs, out_bonds, ref_bonds):
    """Return the Kekule assignment of one ring system that best matches the reference.

    Candidates are all 1/2 assignments of pairs that give every atom the same
    number of double bonds as in the output. The output's own assignment wins ties.
    """
    need, left = defaultdict(int), defaultdict(int)
    for pair in pairs:
        for a in pair:
            need[a] += out_bonds[pair] == 2
            left[a] += 1
    target = [ref_bonds.get(pair, 0) for pair in pairs]
    current = [out_bonds[pair] for pair in pairs]
    best = [sum(a == b for a, b in zip(current, target)), current[:]]

    def search(k, score):
        if score + len(pairs) - k <= best[0]:
            return
        if k == len(pairs):
            best[:] = [score, current[:]]
            return
        i, j = pairs[k]
        left[i] -= 1
        left[j] -= 1
        for order in (2, 1):
            d = order == 2
            if d and not (need[i] and need[j]):
                continue
            need[i] -= d
            need[j] -= d
            if left[i] >= need[i] and left[j] >= need[j]:
                current[k] = order
                search(k + 1, score + (order == target[k]))
            need[i] += d
            need[j] += d
        left[i] += 1
        left[j] += 1

    search(0, 0)
    return dict(zip(pairs, best[1]))


def best_kekule(out_bonds, ref_bonds):
    """Return out_bonds with ring bonds rearranged to the equivalent Kekule structure that best matches ref_bonds.

    Two Kekule structures with the same bonds and the same number of double
    bonds at every atom differ by flipping (1 <-> 2) rings that alternate 1/2,
    so only bonds on rings of single/double bonds can change.
    """
    best = dict(out_bonds)
    for pairs in ring_systems(ring_bonds(out_bonds)):
        best.update(best_system(pairs, out_bonds, ref_bonds))
    return best


def compare_files(output, reference):
    """Return [records, bonds, matched bonds, fully matched molecules]."""
    s = [0, 0, 0, 0]
    pairs = zip_longest(read_records(output), read_records(reference))
    for k, (out, ref) in enumerate(pairs):
        if out is None or ref is None:
            raise ValueError(f"{output} and {reference} differ in record count")
        out_name, out_atoms, out_bonds = parse_bonds(out)
        ref_name, ref_atoms, ref_bonds = parse_bonds(ref)
        if (out_name, out_atoms) != (ref_name, ref_atoms):
            raise ValueError(
                f"record {k + 1}: {out_name} ({out_atoms} atoms) vs "
                f"{ref_name} ({ref_atoms} atoms)"
            )
        out_bonds = best_kekule(out_bonds, ref_bonds)
        pairs_union = out_bonds.keys() | ref_bonds.keys()
        matched = sum(out_bonds.get(p, 0) == ref_bonds.get(p, 0) for p in pairs_union)
        s[0] += 1
        s[1] += len(pairs_union)
        s[2] += matched
        s[3] += matched == len(pairs_union)
    return s


def pct(num, den):
    return 100 * num / den if den else float("nan")


def row(label, n_rec, n_bonds, n_matched, n_exact):
    return (
        f"{label:<24} {n_rec:>8} {n_bonds:>9} {n_matched:>9} "
        f"{pct(n_matched, n_bonds):>7.2f}% {pct(n_exact, n_rec):>7.2f}%"
    )


def main(output, reference):
    stats = compare_files(output, reference)
    header = f"{'':<24} {'records':>8} {'bonds':>9} {'matched':>9} {'bond %':>8} {'mol %':>8}"
    print(header)
    print("-" * len(header))
    print(row("total", *stats))


if __name__ == "__main__":
    parser = ArgumentParser(
        description="Percentage of exact bond type matches between output and reference SDF files."
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=HERE / "output" / "generated.sdf",
        help="generated SDF file",
    )
    parser.add_argument(
        "-r",
        "--reference",
        type=Path,
        default=HERE / "reference" / "reference.sdf",
        help="reference SDF file, record-aligned with --output",
    )
    args = parser.parse_args()
    main(args.output, args.reference)