import pkg_resources
if pkg_resources.get_distribution('CGRtools').version.split('.')[:2] != ['4', '0']:
    print('WARNING. Tutorial was tested on 4.0 version of CGRtools')
else:
    print('Welcome!')

# load data for tutorial
from pickle import load
from traceback import format_exc

with open('molecules.dat', 'rb') as f:
    molecules = load(f) # list of MoleculeContainer objects
with open('reactions.dat', 'rb') as f:
    reactions = load(f) # list of ReactionContainer objects

m1, m2, m3, m4 = molecules # molecule
m7 = m3.copy()
m11 = m3.copy()
m11.standardize()
m7.standardize()
r1 = reactions[0] # reaction
m5 = r1.reactants[0]
m8 = m7.substructure([4, 5, 6, 7, 8, 9])
m10 =  r1.products[0].copy()

"""
Explanation: 1. Data types and operations with them

(c) 2019, 2020 Dr. Ramil Nugmanov;
(c) 2019 Dr. Timur Madzhidov; Ravil Mukhametgaleev

Installation instructions of CGRtools package information and tutorial's files see on https://github.com/stsouko/CGRtools
NOTE: Tutorial should be performed sequentially from the start. Random cell running will lead to unexpected results.
End of explanation
"""


from CGRtools.containers import * # import all containers

"""
Explanation: CGRtools has subpackage containers with data structures classes:

MoleculeContainer - for molecular structure
ReactionContainer - for chemical reaction 
CGRContainer - for Condensed Graph of Reaction
QueryContainer - queries for substructure search in molecules
QueryCGRContainer - queries for substructure search in CGRs
End of explanation
"""


m1.meta # dictionary for molecule properties storage. For example, DTYPE/DATUM fields of SDF file are read into this dictionary

m1 # MoleculeContainer supports depiction and graphic representation in Jupyter notebooks.

m1.depict() # depiction returns SVG image in format string

with open('molecule.svg', 'w') as f: # saving image to SVG file
    f.write(m1.depict())

m_copy = m1.copy() # copy of molecule
m_copy

len(m1) # get number of atoms in molecule
# or 
m1.atoms_count

m1.bonds_count # number of bonds

m1.atoms_numbers # list of atoms numbers

"""
Explanation: 1.1. MoleculeContainer
Molecules are represented as undirected graphs. Molecules contain Atom objects and Bond objects.
Atom objects are represented as dictionary with unique number for each atom as key.  
Bond objects are stored as sparse matrix with adjacent atoms pair as keys for rows and columns.
Hereafter, atom number is unique integer used to enumerate atoms in molecule. Please, don't confuse it with element number in Periodic Table, hereafter called element number.
Methods for molecule handling and the arguments of MoleculeContainer are described below.
End of explanation
"""


# iterate over atoms using its numbers
list(m1.atoms())  # works the same as dict.items()

# iterate over bonds using adjacent atoms numbers
list(m1.bonds())

# access to atom by number
m1.atom(1)

try:
    m1.atom(10) # raise error for absent atom numbers
except KeyError:
    print(format_exc())

# access to bond using adjacent atoms numbers
m1.bond(1, 4)

try:
    m1.bond(1, 3) # raise error for absent bond
except KeyError:
    print(format_exc())

"""
Explanation: Each structure has additional atoms attributes: number of neighbors and hybridization.
The following notations are used for hybridization of atoms. Values are given as numbers below (in parenthesis symbols that are used in SMILES-like signatures are shown):

1 (s) - all bonds of atom are single, i.e. sp3 hybridization
2 (d) - atom has one double bond and others are single, i.e. sp2 hybridization
3 (t) - atom has one triple or two double bonds and other are single, i.e. sp hybridization
4 (a) - atom is in aromatic ring

Neighbors and hybridizations atom attributes are required for substructure operations and structure standardization. See below
End of explanation
"""


a = m1.atom(1)

# access to information
a.atomic_symbol # element symbol

a.charge # formal charge

a.is_radical # atom radical state

a.isotope # atom isotope. Default isotope if not set. Default isotopes are the same as used in InChI notation

a.x # coordinates
a.y
#or 
a.xy

a.neighbors # Number of neighboring atoms. It is read-only.

a.hybridization # Atoms hybridization. It is read-only.

try:
    a.hybridization = 2 # Not assignable. Read-only! Thus error is raised.
except AttributeError:
    print(format_exc())

"""
Explanation: Atom objects are dataclasses which store information about:

element
isotope
charge
radical state
xy coordinates

Also atoms has methods for data integrity checks and include some internally used data.
End of explanation
"""


a.charge = 1
m1

a.charge = 0
a.is_radical = True
m1

# bond objects also are data-like classes which store information about bond order
b = m1.bond(3, 4)
b.order

try:
    b.order = 1 # order change not possible
except AttributeError:
    print(format_exc())

"""
Explanation: Atomic attributes are assignable.
CGRtools has integrity checks for verification of changes induced by user
End of explanation
"""


m1.delete_bond(3, 4)
m1

"""
Explanation: Bonds are Read-only
For bond modification one should to use delete_bond method to break bond and add_bond for creating new.
End of explanation
"""


m1.delete_atom(3)
m1

m_copy # copy unchanged!

"""
Explanation: Method delete_atom removes atom from the molecule
End of explanation
"""


int(a)
# 61705 == 000001111 0001000 0100 1
# 000001111 == 15 isotope
# 0001000 == 8 Oxygen
# 0100 == 4 (4 - 4 = 0) uncharged
# 1 == 1 is radical

int(b)  # bonds are encoded by their order

a = m_copy.atom(1)
print(a.implicit_hydrogens) # get number of implicit hydrogens on atom 1
print(a.explicit_hydrogens) # get number of explicit hydrogens on atom 1
print(a.total_hydrogens) # get total number of hydrogens on atom 1

m1

m1.check_valence() # return list of numbers of atoms with invalid valences

m4 # molecule with valence errors

m4.check_valence()

m3

m3.sssr # Method for application of Smallest Set of Smallest Rings algorithm for rings 
        # identification. Returns tuple of tuples of atoms forming smallest rings

"""
Explanation: Atoms and bonds objects can be converted into integer representation that could be used to classify their types.
Atom type is represented by 21 bit code rounded to 32 bit integer number:

9 bits are used for isotope (511 posibilities, highest known isotope is ~300)
7 bits stand for atom number (2 ** 7 - 1 == 127, currently 118 elements are presented in Periodic Table)
4 bits stand for formal charge. Charges range from -4 to +4 rescaled to range 0-8
1 bit are used for radical state.
End of explanation
"""


m2 # it's a salt represented as one graph

m2.connected_components # tuple of tuples of atoms belonging to graph components

anion, cation = m2.split() # split molecule to components

anion # graph of only one salt component

cation # graph of only one salt component

"""
Explanation: Connected components.
Sometimes molecules has disconnected components (salts etc).
One can find them and split molecule to separate components.
End of explanation
"""


salt = anion | cation 
# or 
salt = anion.union(cation)
salt # this graph has disconnected components, it is considered single compound now

"""
Explanation: Union of molecules
Sometimes it is more convenient to represent salts as ion pair. Otherwise ambiguity could be introduced, for example in reaction of salt exchange:
Ag+ + NO3- + Na+ + Br- = Ag+ + Br- + Na+ + NO3-. Reactants and products sets are the same. 
In this case one can combine anion-cation pair into single graph. It could be convenient way to represent other molecule mixtures.
End of explanation
"""


sub = m3.substructure([4,5,6,7,8,9])  # substructure with passed atoms
sub

"""
Explanation: Substructures could be extracted from molecules.
End of explanation
"""


aug = m3.augmented_substructure([10], deep=2) #  atom 10 is Nitrogen
aug

"""
Explanation: augmented_substructure is a substructure consisting from atoms and a given number of shells of neighboring atoms around it.
deep argument is a number of considered shells. 
It also returns projection by default.
End of explanation
"""


m5.atoms_order

"""
Explanation: Atoms Ordering.
This functionality is used for canonic numbering of atoms in molecules. Morgan algorithm is used for atom ranking. Property atoms_order returns dictionary of atom numbers as keys and their ranks according to canonicalization as values. Equal rank mean that atoms are symmetric (are mapped to each other in automorhisms).
End of explanation
"""


m5

remapped = m5.remap({4:2}, copy=True)
remapped

"""
Explanation: Atom number can be changed by remap method.
This method is useful when it is needed to change order of atoms in molecules. First argument to remap method is dictionary with existing atom numbers as keys and desired atom number as values. It is possible to change atom numbers for only part of atoms. Atom numbers could be non-sequencial but need to be unique. 
If argument copy is set True new object will be created, else existing molecule will be changed. Default is False.
End of explanation
"""


r1 # depiction supported

r1.meta

print(r1.reactants, r1.products)  # Access to lists of reactant and products.
reactant1, reactant2, reactant3 = r1.reactants
product = r1.products[0]

"""
Explanation: 1.2. ReactionContainer
ReactionContainer objects has the following properties:

reactants - list of reactants molecules
reagents - list of reagents molecules
products - list of products molecules
meta - dictinary of reaction metadata (DTYPE/DATUM block in RDF)
End of explanation
"""


cgr1 = m7 ^ m8 # CGR from molecules
# or 
cgr1 = m7.compose(m8)
print(cgr1)
cgr1

r1

cgr2 = ~r1 # CGR from reactions
# or 
cgr2 = r1.compose()
print(cgr2) # signature is printed out.
cgr2.clean2d()
cgr2

a = cgr2.atom(2) # atom access is the same as for MoleculeContainer

a.atomic_symbol # element attribute

a.isotope # isotope attribute

"""
Explanation: Reactions also has standardize, kekule, thiele, implicify_hydrogens, explicify_hydrogens, etc methods (see part 3). These methods are applied independently to every molecule in reaction.
1.3. CGR
CGRContainer object is similar to MoleculeConrtainer, except some methods. The following methods are not suppoted for CGRContainer:

standardization methods
hydrogens count methods
check_valence

CGRContainer also has some methods absent in MoleculeConrtainer:

centers_list
center_atoms
center_bonds

CGRContainer is undirected graph. Atoms and bonds in CGR has two states: reactant and product.
Composing to CGR
As mentioned above, atoms in MoleculeContainer have unique numbers. These numbers are used as atom-to-atom mapping in CGRtools upon CGR creation. Thus, atom order for molecules in reaction should correspond to atom-to-atom mapping.  
Pair of molecules can be transformed into CGR. Notice that, the same atom numbers in reagents and products imply the same atoms.
Reaction also can be composed into CGR. Atom numbers of molecules in reaction are used as atom-to-atom mapping of reactants to products.
End of explanation
"""


a.charge # charge of atom in reactant

a.p_charge # charge of atom in product

a.p_is_radical # radical state of atom in product.

a.neighbors # number of neighbors of atom in reactant

a.p_neighbors # number of neighbors of atom in product

a.hybridization # hybridization of atom in reactant. 1 means only single bonds are incident to atom

a.p_hybridization # hybridization of atom in product. 1 means only single bonds are incident to atom

b = cgr1.bond(4, 10) # take bond

"""
Explanation: For CGRContainer attributes charge, is_radical, neighbors and hybridization refer to atom state in reactant of reaction;  arguments p_charge, p_is_radical, p_neighbors and p_hybridization could be used to extract atom state in product part in reaction.
End of explanation
"""


b.order # bond order in reactant

b.p_order is None # bond order in product in None

"""
Explanation: Bonds has order and p_order attribute
If order attribute value is None, it means that bond was formed
If p_order is None, it means that bond was broken  
Both order and p_order can't be None
End of explanation
"""


reactant_part, product_part = ~cgr1 # CGR of unbalanced reaction is decomposed back into reaction
# or 
reactant_part, product_part = cgr1.decompose()

reactant_part # reactants extracted. One can notice it is initial molecule

product_part #extracted products. Originally benzene was the product.

"""
Explanation: CGR can be decomposed back to reaction, i.e. reactants and products.
Notice that CGR can lose information in case of unbalanced reactions (where some atoms of reactant does not have counterpart in product, and vice versa). Decomposition of CGRs for unbalanced reactions back to reaction may lead to strange (and erroneous) structures.
End of explanation
"""


decomposed = ReactionContainer.from_cgr(cgr2)
decomposed.clean2d()
decomposed

"""
Explanation: For decomposition of CGRContainer back into ReactionContainer ReactionContainer.from_cgr constructor method can be used.
End of explanation
"""


r1 # compare with initial reaction

"""
Explanation: You can see that water absent in products initially was restored. 
This is a side-effect of CGR decomposing that could help with reaction balancing. 
But balancing using CGR decomposition works correctly only if minor part atoms are lost 
but multiplicity and formal charge are saved. In next release electronic state balansing will be added.
End of explanation
"""


from CGRtools.containers import*

m10 # ether

carb = m10.substructure([5,7,2], as_query=True) # extract of carboxyl fragment
print(carb)
carb

"""
Explanation: 1.4 Queries
CGRtools supports special objects for Queries. Queries are designed for substructure isomorphism. User can set number of neighbors and hybridization by himself (in molecules they could be calculated but could not be changed).
Queries don't have reset_query_marks method
End of explanation
"""


cgr_q = cgr1.substructure(cgr1, as_query=True) # transfrom CGRContainer into QueryCGRContainer
#or
cgr_q = QueryCGRContainer() | cgr1 # Union of Query container with CGR or Molecule gives QueryCGRContainer
print(cgr_q)                    # print out signature of query
cgr_q

"""
Explanation: CGRs also can be transformed into Query.
QueryCGRContainer is similar to QueryContainer class for CGRs and has the same API.
QueryCGRContainer take into account state of atoms and bonds in reactant and product, including neighbors and hybridization
End of explanation
"""


from CGRtools.containers import MoleculeContainer
from CGRtools.containers.bonds import Bond
from CGRtools.periodictable import Na

m = MoleculeContainer() # new empty molecule

m.add_atom('C')  # add Carbon atom using element symbol
m.add_atom(6)    # add Carbon atom using element number. {'element': 6} is not valid, but {'element': 'O'} is also acceptable
m.add_atom('O', charge=-1) # add negatively charged Oxygen atom. Similarly other atomic properties can be set

# add_atom has second argument for setting atom number. 
# If not set, the next integer after the biggest among already created will be used.
m.add_atom(Na(23), 4, charge=1) # For isotopes required element object construction.

m.add_bond(1, 2, 1) # add bond with order = 1 between atoms 1 and 2
m.add_bond(3, 2, Bond(1)) # the other possibility to set bond order

m.clean2d() #experimental function to calculate atom coordinates. Has number of flaws yet
m

"""
Explanation: 1.5. Molecules, CGRs, Reactions construction
CGRtools has API for objects construction from scratch.
CGR and Molecule has methods add_atom and add_bond for adding atoms and bonds.
End of explanation
"""


r = ReactionContainer(reactants=[m1], products=[m11]) # one-step way to construct reaction
# or
r = ReactionContainer([m1], [m11]) # first list of MoleculeContainers is interpreted as reactants, second one - as products

r

r.fix_positions() # this method fixes coordinates of molecules in reaction without calculation of atoms coordinates.
r

"""
Explanation: Reactions can be constructed from molecules.
Reactions are tuple-like objects. Modification impossible.
End of explanation
"""


q = QueryContainer() # creation of empty container
q.add_atom('N') # add N atom, any isotope, not radical, neutral charge, 
                # number of neighbors and hybridization are irrelevant
q.add_atom('C', neighbors=[2, 3], hybridization=2) # add carbon atom, any isotope, not radical, neutral charge, 
                # has 2 or 3 explicit neighbors and sp2 hybridization
q.add_atom('O', neighbors=1)
q.add_bond(1, 2, 1) # add single bond between atom 1 and 2 
q.add_bond(2, 3, 2) # add double bond between atom 1 and 2 
# any amide group will fit this query

print(q) # print out signature (SMILES-like)
q.clean2d()
q

"""
Explanation: QueryContainers can be constructed in the same way as MoleculeContainers.
Unlike other containers QueryContainers additionally support atoms, neighbors and hybridization lists.
End of explanation
"""


from CGRtools.periodictable import Core, C, O

class Marked(Core):
    __slots__ = '__mark' # all new attributes should be slotted!

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.__mark = None # set default value for added attribute

    @property
    def mark(self):        # created new property 
        return self.__mark

    @mark.setter
    def mark(self, mark):
        # do some checks and calculations
        self.__mark = mark
    
    def __repr__(self):
        if self.__isotope:
            return f'{self.__class__.__name__[6:]}({self.__isotope})'
        return f'{self.__class__.__name__[6:]}()'
    
    @property
    def atomic_symbol(self) -> str:
        return self.__class__.__name__[6:]


class MarkedC(Marked, C):
    pass


class MarkedO(Marked, O):
    pass

m = MoleculeContainer() # create newly developed container MarkedMoleculeContainer
m.add_atom(MarkedC())   # add custom atom C
m.add_atom(MarkedO())   # add custom atom O
m.add_bond(1, 2, 1)

m.atom(2).mark = 1  # set mark on atom.

print(m)
m.clean2d()
m

m.atom(2).mark # one can return mark

"""
Explanation: 1.6. Extending CGRtools
You can easily customize CGRtools for your tasks.
CGRtools is OOP-oriented library with subclassing and inheritance support.
As an example, we show how special marks on atoms for ligand donor centers can be added.
End of explanation
"""


