import numpy as np
from rdkit import Chem

# examples
# C=[NH+]C
component = "C=[NH+]"
component = "[NH+]C"

# C[S+](C)C
component = "C[S+]"


print(component)

mc = Chem.MolFromSmiles(component)
n_atoms = mc.GetNumAtoms()
n_bonds = len(mc.GetBonds())

charges = np.zeros(n_atoms, dtype=int)

for idx in range(n_atoms):
    atom = mc.GetAtomWithIdx(idx)
    atom.SetNumExplicitHs(0)
    charge = atom.GetFormalCharge()
    charges[idx] = charge
    atom.SetFormalCharge(0)

component = Chem.MolToSmiles(mc, canonical=False)
component = component.replace("[", "").replace("]","")

mc = Chem.MolFromSmiles(component)

for idx, charge in zip(range(n_atoms), charges):
    atom = mc.GetAtomWithIdx(idx)
    atom.SetFormalCharge(charge)

component = Chem.MolToSmiles(mc)
print(component)