"""
mermaid_to_psmiles.py

Parse a polymer Mermaid MoleCode graph (as generated by polymer_to_mermaid.py)
back into a PSMILES repeat-unit string with two * attachment points.

The MoleCode format contract (from polymer_to_mermaid.py):
  - Nodes:   {prefix}_{Symbol}{n}[{label}]   e.g.  B0_C1[CH2]
  - Bonds:   "---" SINGLE | "===" DOUBLE | "-.-" TRIPLE | "<-->" AROMATIC
             "===|E|" / "===|Z|" E/Z-stereo double
  - Attachment points identified by:
             TL --- {entry_node}   (left  * )
             {exit_node} --- TR    (right * )

Design:
  - Only the element symbol is read from node labels; H counts are inferred
    by RDKit SanitizeMol from connectivity and valence.  This means the model
    does NOT need to update H-count labels when adding/removing branches.
  - Aromatic bonds (<-->) are passed as BondType.AROMATIC; SanitizeMol
    perceives aromaticity and sets atom/bond flags automatically.

Round-trip test (all 50 editing task inputs should recover exactly):
    python3 mermaid_to_psmiles.py
"""

from __future__ import annotations

import re
import sys
import json
from pathlib import Path
from typing import Optional

from rdkit import Chem
from rdkit.Chem import RWMol


# ── Label → element symbol ────────────────────────────────────────────────────

# Two-letter symbols that must be matched before single-letter
_TWO_LETTER = frozenset(["Si", "Cl", "Br", "Se", "Te", "Mg", "Zn", "Fe", "Cu"])

def parse_element(label: str) -> tuple[str, int]:
    """
    Extract element symbol and formal charge from a node label.
    H counts in the label are intentionally ignored — RDKit infers them.

    Examples:
        "CH2"  → ("C",  0)
        "NH"   → ("N",  0)
        "O"    → ("O",  0)
        "N(+)" → ("N", +1)
        "Si"   → ("Si", 0)
    """
    label = label.strip()

    # Formal charge
    charge = 0
    m = re.search(r'\((\d*)([+\-])\)', label)
    if m:
        mag = int(m.group(1)) if m.group(1) else 1
        charge = mag if m.group(2) == "+" else -mag
        label = label[:m.start()]

    # Two-letter element symbols first
    for sym in _TWO_LETTER:
        if label.startswith(sym):
            return sym, charge

    # Single-letter
    if label and label[0].isupper():
        return label[0], charge

    raise ValueError(f"Cannot parse element from label: {label!r}")


# ── Bond symbol → RDKit BondType ──────────────────────────────────────────────

_BOND_MAP: dict[str, Chem.BondType] = {
    "---":     Chem.BondType.SINGLE,
    "===":     Chem.BondType.DOUBLE,
    "===|E|":  Chem.BondType.DOUBLE,
    "===|Z|":  Chem.BondType.DOUBLE,
    "-.-":     Chem.BondType.TRIPLE,
    "<-->":    Chem.BondType.AROMATIC,
}

_EDGE_RE = re.compile(
    r'^\s+(\w+)\s+'
    r'(-{3}|={3}(?:\|[EZ]\|)?|\-\.\-|<-->)'
    r'\s+(\w+)\s*$'
)

_NODE_RE = re.compile(r'^\s+(\w+)\[([^\]]+)\]\s*$')

# Absolute CIP chirality suffix on an atom id, e.g. ``B0_C2_R`` → "R".
# The preceding char is always a digit (the per-element counter), so this never
# collides with an element symbol such as sulfur (``B0_S1``).
_CHIRAL_RE = re.compile(r'\d_([RS])$')


# ── Stereochemistry restoration (after sanitize / full topology) ──────────────

def _restore_chirality(mol: RWMol, idx_to_cip: dict[int, str]) -> None:
    """Restore absolute CIP R/S by brute-forcing the chiral tag.

    CW/CCW is order-dependent, so we try each tag, recompute CIP, and keep the
    one whose ``_CIPCode`` matches the desired absolute label (mirrors the
    small-molecule converter).
    """
    if not idx_to_cip:
        return
    for idx, desired in idx_to_cip.items():
        atom = mol.GetAtomWithIdx(idx)
        matched = False
        for tag in (Chem.ChiralType.CHI_TETRAHEDRAL_CW,
                    Chem.ChiralType.CHI_TETRAHEDRAL_CCW):
            atom.SetChiralTag(tag)
            try:
                Chem.AssignStereochemistry(mol, cleanIt=True, force=True)
            except Exception:
                continue
            if atom.HasProp('_CIPCode') and atom.GetProp('_CIPCode') == desired:
                matched = True
                break
        if not matched:
            atom.SetChiralTag(Chem.ChiralType.CHI_UNSPECIFIED)
    try:
        Chem.AssignStereochemistry(mol, cleanIt=False, force=True)
    except Exception:
        pass


def _restore_double_bond_stereo(mol: RWMol, records: list) -> None:
    """Restore ===|E| / ===|Z| double-bond configuration.

    Sets the double-bond stereo + reference atoms, then derives the neighbouring
    single-bond directions so the SMILES writer actually emits ``/`` and ``\``.
    """
    applied = False
    for idx1, idx2, stereo in records:
        bond = mol.GetBondBetweenAtoms(idx1, idx2)
        if bond is None:
            continue
        n1 = [n.GetIdx() for n in mol.GetAtomWithIdx(idx1).GetNeighbors()
              if n.GetIdx() != idx2]
        n2 = [n.GetIdx() for n in mol.GetAtomWithIdx(idx2).GetNeighbors()
              if n.GetIdx() != idx1]
        if n1 and n2:
            bond.SetStereoAtoms(n1[0], n2[0])
            bond.SetStereo(Chem.BondStereo.STEREOE if stereo == 'E'
                           else Chem.BondStereo.STEREOZ)
            applied = True
    if applied:
        try:
            Chem.SetDoubleBondNeighborDirections(mol)
        except Exception:
            pass


# ── Main converter ────────────────────────────────────────────────────────────

def mermaid_to_psmiles(mermaid: str) -> Optional[str]:
    """
    Convert a polymer Mermaid MoleCode string back to a PSMILES repeat unit.

    Returns canonical PSMILES (with two * attachment points) or None if
    parsing/sanitization fails.
    """
    lines = mermaid.splitlines()

    # 1. Collect node declarations (skip TL / TR terminus nodes)
    nodes: dict[str, tuple[str, int]] = {}  # nid → (element, formal_charge)
    chirality: dict[str, str] = {}          # nid → "R" / "S"
    for line in lines:
        m = _NODE_RE.match(line)
        if m:
            nid, label = m.group(1), m.group(2)
            if nid in ("TL", "TR"):
                continue
            try:
                symbol, charge = parse_element(label)
                nodes[nid] = (symbol, charge)
            except ValueError:
                continue
            cm = _CHIRAL_RE.search(nid)
            if cm:
                chirality[nid] = cm.group(1)

    # 2. Collect edges and identify attachment nodes
    edges: list[tuple[str, Chem.BondType, str, Optional[str]]] = []
    entry_nodes: set[str] = set()
    exit_nodes:  set[str] = set()

    for line in lines:
        m = _EDGE_RE.match(line)
        if not m:
            continue
        a, bond_sym, b = m.group(1), m.group(2), m.group(3)
        bt = _BOND_MAP.get(bond_sym, Chem.BondType.SINGLE)
        ez = 'E' if bond_sym == '===|E|' else ('Z' if bond_sym == '===|Z|' else None)

        if a == "TL":
            entry_nodes.add(b)
        elif b == "TL":
            entry_nodes.add(a)
        elif b == "TR":
            exit_nodes.add(a)
        elif a == "TR":
            exit_nodes.add(b)
        else:
            if a in nodes and b in nodes:
                edges.append((a, bt, b, ez))

    if not entry_nodes or not exit_nodes:
        return None

    entry_nid = next(iter(entry_nodes))
    exit_nid  = next(iter(exit_nodes))

    # 3. Build RDKit RWMol — H counts inferred by SanitizeMol
    rw = RWMol()
    nid_to_idx: dict[str, int] = {}

    for nid, (symbol, charge) in nodes.items():
        try:
            atom = Chem.Atom(symbol)
        except Exception:
            return None
        atom.SetFormalCharge(charge)
        # No explicit H — let RDKit compute from valence after sanitization
        nid_to_idx[nid] = rw.AddAtom(atom)

    # Wildcard * atoms for attachment points
    ap_left  = rw.AddAtom(Chem.Atom(0))
    ap_right = rw.AddAtom(Chem.Atom(0))

    stereo_records: list[tuple[int, int, str]] = []  # (idx1, idx2, 'E'/'Z')
    for a, bt, b, ez in edges:
        if a in nid_to_idx and b in nid_to_idx:
            try:
                rw.AddBond(nid_to_idx[a], nid_to_idx[b], bt)
            except Exception:
                return None
            if ez:
                stereo_records.append((nid_to_idx[a], nid_to_idx[b], ez))

    if entry_nid not in nid_to_idx or exit_nid not in nid_to_idx:
        return None
    rw.AddBond(ap_left,  nid_to_idx[entry_nid], Chem.BondType.SINGLE)
    rw.AddBond(nid_to_idx[exit_nid], ap_right,  Chem.BondType.SINGLE)

    # 4. Sanitize
    try:
        Chem.SanitizeMol(rw)
    except Exception:
        try:
            Chem.SanitizeMol(
                rw,
                Chem.SanitizeFlags.SANITIZE_ALL
                ^ Chem.SanitizeFlags.SANITIZE_PROPERTIES,
            )
        except Exception:
            return None

    # 5. Restore stereochemistry (R/S and E/Z) after the full topology exists.
    _restore_chirality(
        rw, {nid_to_idx[n]: cip for n, cip in chirality.items() if n in nid_to_idx}
    )
    _restore_double_bond_stereo(rw, stereo_records)

    return Chem.MolToSmiles(rw)


# ── Round-trip test ───────────────────────────────────────────────────────────

def round_trip_test(questions_path: str = "chemiq_results/polymer_editing_questions.jsonl"):
    """
    Verify mermaid_to_psmiles(mermaid_graph) == canonical(input_smiles)
    for all editing task inputs.
    """
    def canon(smi):
        m = Chem.MolFromSmiles(smi)
        return Chem.MolToSmiles(m) if m else None

    records = [json.loads(l) for l in open(questions_path) if l.strip()]
    ok = fail = 0
    for rec in records:
        expected = canon(rec["input_repeat_smiles"])
        recovered = mermaid_to_psmiles(rec["mermaid_graph"])
        if recovered and canon(recovered) == expected:
            ok += 1
        else:
            fail += 1
            print(f"  FAIL  {rec['name']}")
            print(f"    original : {rec['input_repeat_smiles']}  →  {expected}")
            print(f"    recovered: {recovered}")

    print(f"\nRound-trip: {ok}/{ok+fail} passed  ({fail} failures)")
    return fail == 0


if __name__ == "__main__":
    success = round_trip_test()
    sys.exit(0 if success else 1)

