from enum import Enum
from functools import cmp_to_key, partial
from itertools import groupby, chain

from indigo import Indigo, IndigoObject
from indigo.renderer import IndigoRenderer

_indigo = Indigo()

renderer = IndigoRenderer(_indigo)
# _indigo.setOption('render-atom-ids-visible', 'true')
_indigo.setOption('render-stereo-style', 'none')
_indigo.setOption('render-superatom-mode', 'collapse')
_indigo.setOption('render-label-mode', 'terminal-hetero')

# smiles = "Cl[C@@H]1[C@@H](N)CCCC1"
smiles = "C[C@@H](C[C@@H](C)Br)C[C@H](CBr)Cl"

m = _indigo.loadMolecule(smiles)


class RSStereoType(str, Enum):
    R = 'R'
    S = 'S'


def stereo_compare(atom1: IndigoObject, atom2: IndigoObject, *, root):
    if (n1:=atom1.atomicNumber()) != (n2:=atom2.atomicNumber()):
        return (n1 > n2) - (n1 < n2)
    atom1_neighbors = tuple(atom for atom in atom1.iterateNeighbors() if atom.index() != root.index())
    atom2_neighbors = tuple(atom for atom in atom2.iterateNeighbors() if atom.index() != root.index())
    if not atom1_neighbors and not atom2_neighbors:  # should be invalid
        return 1
    if not atom1_neighbors:
        return -1
    if not atom2_neighbors:
        return 1
    left = max(atom1_neighbors, key=cmp_to_key(partial(stereo_compare, root=atom1)))
    right = max(atom2_neighbors, key=cmp_to_key(partial(stereo_compare, root=atom2)))
    return stereo_compare(left, right, root=root)


def is_clockwise(p1, p2, p3):
    if((p2[0]-p1[0])*(p3[1]-p1[1]) - (p2[1]-p1[1])*(p3[0]-p1[0]))<0:
        return True
    return False


def mark_stereochemistry(m: IndigoObject):
    assert m.dbgInternalType() == '#02: <molecule>'  # add reaction supports?
    m.layout()
    # 获取所有手性原子
    chiral_centers = tuple(m.iterateStereocenters())
    # 所有手性中心原子的下标集合
    chiral_centers_idx = {atom.index() for atom in chiral_centers}
    # 查出与手性中心相连的所有键
    related_bonds = sorted(tuple(chain.from_iterable(
        ((bond.source().index(), bond), (bond.destination().index(), bond),)
        for bond in m.iterateBonds()
        if bond.source().index() in chiral_centers_idx
        or bond.destination().index() in chiral_centers_idx
    )), key=lambda x: x[0])
    chiral_centers_mapping = {
        k: tuple(x[1] for x in g)
        for k, g in groupby(related_bonds, key=lambda x: x[0])
        if k in chiral_centers_idx
    }
    # 判断最终结果是否需要反转
    chiral_centers_type = {
        idx: False
        if any(b.bondStereo()==5 for b in bonds)
        else (True if any(b.bondStereo()==6 for b in bonds) else None)
        for idx, bonds in chiral_centers_mapping.items()
    }

    for chiral_center in chiral_centers:
        center_idx = chiral_center.index()
        if (should_flip := chiral_centers_type[center_idx]) is None:
            continue
        # TODO skip chiral_center is not C?
        neighbors = tuple(atom for atom in chiral_center.iterateNeighbors())  # TODO
        # skip invalid stereo center
        assert len(neighbors) >= 3
        neighbors = sorted(
            neighbors, key=cmp_to_key(partial(stereo_compare, root=chiral_center)),reverse=True
        )
        # TODO assert chiral_center inside of area of neighbors
        atom1, atom2, atom3, *_ = neighbors
        clockwise = is_clockwise(atom1.xyz(), atom2.xyz(), atom3.xyz())
        # consider bond type flip
        if should_flip:
            clockwise = not clockwise
        if clockwise:
            stereo_type = RSStereoType.R
        else:
            stereo_type = RSStereoType.S
        # 根据手性，设定DataSGroup
        m.addDataSGroup([center_idx], [], 'stereochemistry', f"({stereo_type})")


mark_stereochemistry(m)
renderer.renderToFile(m, "D:/asdf.png")

if __name__ == '__main__':
    pass