#!/usr/bin/env python
# coding: utf-8

# ## Create the systems
#
# ### Create the ring molecule from SMILES
#
# First, we create ``ring_mol`` from the SMILES code, and then we cut the O-C and N-C bonds and increase the distance between the fragments. We need to first cut the bonds before increasing the distance, since ``ChemicalSystem.set_distance`` requires separate molecule fragments to work on.

from scm.plams import view
from scm.base import ChemicalSystem

# from_smiles creates bonds between atoms
ring_mol = ChemicalSystem.from_smiles("C1CON=N1")
view(
    ring_mol,
    width=350,
    height=350,
    direction="tilt_pca3",
    show_atom_labels=True,
    atom_label_type="Name",
)


# ### Identify which bonds to cut
#
# It's easiest if you know the indices of the bonded atoms. Then you can just cut those bonds, e.g. ``ring_mol.bonds.remove_bonds_between_atoms(0, 4)``. That would remove the N-C bond in the above picture (there the indices start with 1, but when calling Python methods they start from 0)
#
# That is fine and the most straightforward way to cut the bonds. However, it is difficult to automate because the atomic indices may change for different structures.
#
# Here we show a way to programmatically identify "an N atom that is bonded to another N and C", and "an O atom that is bonded to an N and C". Then we cut those corresponding bonds.

from typing import List, Tuple, Iterable, Optional


def find_first_matching(
    cs: ChemicalSystem, elem: str, neighbors: Iterable[str]
) -> Optional[Tuple[int, List[int], List[str]]]:
    """Finds the first atom with element ``elem`` and neighbors ``neighbors``.

    Returns a 3-tuple: (index_of_atom, list_of_neighbor_indices, list_of_neighbor_elements)

    In the return value the list_of_neighbor_indices and list_of_neighbor_elements
    have the same length and correspond to each other.
    """

    sorted_target_neighbors = sorted(neighbors)
    for i, at in enumerate(cs):
        if at.symbol != elem:
            continue
        my_neighbors = list(cs.bonds.get_bonded_atoms(i))
        my_neighbors_elems = [cs.atoms[j].symbol for j in my_neighbors]
        if sorted(my_neighbors_elems) == sorted_target_neighbors:
            return i, my_neighbors, my_neighbors_elems

    return None


sep_mol = ring_mol.copy()

# N bonded to N and C
iN, N_neighs, N_neighs_elems = find_first_matching(sep_mol, "N", ["N", "C"])
iN_C = N_neighs[N_neighs_elems.index("C")]  # index of neighboring C atom

# O bonded to N and C
iO, O_neighs, O_neighs_elems = find_first_matching(sep_mol, "O", ["N", "C"])
iO_C = O_neighs[O_neighs_elems.index("C")]  # index of neighboring C atom


print(f"I will cut the bonds between atoms {iN}-{iN_C} and {iO}-{iO_C} (zero-based indexing)")


# ### Cut the bonds

sep_mol.bonds.remove_bonds_between_atoms(iN, iN_C)
sep_mol.bonds.remove_bonds_between_atoms(iO, iO_C)
view(sep_mol, width=150, height=150, direction="tilt_pca3", picture_path="picture1.png")


# ### Increase the distance
#
# Now increase the distance. When you change the distance, all distances within the separate molecules stay the same.

sep_mol.set_distance(iN, iN_C, 2.5)  # angstrom
view(sep_mol, width=150, height=150, direction="tilt_pca3", picture_path="picture2.png")


# ### Preoptimize the separated and ring systems

from scm.plams import Settings, AMSJob, view

preopt_s = Settings()
preopt_s.input.ams.Task = "GeometryOptimization"
preopt_s.input.DFTB.Model = "GFN1-xTB"
preopt_ring_job = AMSJob(settings=preopt_s, molecule=ring_mol, name="preopt_ring_mol")
preopt_ring_job.run()
preopt_sep_job = AMSJob(settings=preopt_s, molecule=sep_mol, name="preopt_sep_mol")
preopt_sep_job.run()
# ## Final molecule results

opt_ring_mol = preopt_ring_job.results.get_main_system()
view(opt_ring_mol, width=150, height=150, direction="tilt_pca3", picture_path="picture3.png")


opt_sep_mol = preopt_sep_job.results.get_main_system()
view(opt_sep_mol, width=150, height=150, direction="tilt_pca3", picture_path="picture4.png")


# ## Copy-pasteable system blocks
#
# When running calculations, bonds are really only used in ForceField calculations. Otherwise, you do not need any bonding information (it is not used). So it can be a bit misleading to store it or print it. Therefore, we here remove all the bonds before printing.

opt_ring_mol.bonds.clear_bonds()
print(opt_ring_mol)


opt_sep_mol.bonds.clear_bonds()
print(opt_sep_mol)