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

################################################################################
#
#   ChemPy - A chemistry toolkit for Python
#
#   Copyright (c) 2012 by Joshua W. Allen (jwallen@mit.edu)
#
#   Permission is hereby granted, free of charge, to any person obtaining a 
#   copy of this software and associated documentation files (the "Software"), 
#   to deal in the Software without restriction, including without limitation
#   the rights to use, copy, modify, merge, publish, distribute, sublicense, 
#   and/or sell copies of the Software, and to permit persons to whom the 
#   Software is furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
#   THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 
#   DEALINGS IN THE SOFTWARE. 
#
################################################################################

"""
This module contains functionality for reading from and writing to the
adjacency list format used by Reaction Mechanism Generator (RMG).
"""
import logging
import re
from .molecule import Atom, Bond, getAtomType
from .group import GroupAtom, GroupBond
from .element import getElement, PeriodicSystem
from rmgpy.exceptions import InvalidAdjacencyListError

class Saturator(object):
    @staticmethod
    def saturate(atoms):
            '''
            Returns a list of atoms that is extended 
            (and bond attributes) by saturating the valency of the non-hydrogen atoms with an 
            appropriate number of hydrogen atoms.
            
            The required number of hydrogen atoms per heavy atom is determined as follows:
            H's =     max number of valence electrons - atom.radicalElectrons
                        - 2* atom.lonePairs - order - atom.charge
            
            '''
            newAtoms = []
            for atom in atoms:
                try:
                    max_number_of_valence_electrons = PeriodicSystem.valence_electrons[atom.symbol]
                except KeyError:
                    raise InvalidAdjacencyListError('Cannot add hydrogens to adjacency list: Unknown orbital for atom "{0}".'.format(atom.symbol))
                
                order = atom.getBondOrdersForAtom()
                    
                number_of_H_to_be_added = max_number_of_valence_electrons - atom.radicalElectrons - 2* atom.lonePairs - int(order) - atom.charge
                
                if number_of_H_to_be_added < 0:
                    raise InvalidAdjacencyListError('Incorrect electron configuration on atom.')
                    
                for _ in range(number_of_H_to_be_added):
                    a = Atom(element='H', radicalElectrons=0, charge=0, label='', lonePairs=0)
                    b = Bond(atom, a, 'S')
                    newAtoms.append(a)
                    atom.bonds[a] = b
                    a.bonds[atom] = b
            atoms.extend(newAtoms)  

class ConsistencyChecker(object):
        
    @staticmethod
    def check_partial_charge(atom):
            '''
            Checks whether the partial charge attribute of the atom checks out with 
            the theoretical one:
            
            '''
            valence = PeriodicSystem.valence_electrons[atom.symbol]
            order = atom.getBondOrdersForAtom()
                
            theoretical = valence - order - atom.radicalElectrons - 2*atom.lonePairs

            if atom.charge != theoretical:
                raise InvalidAdjacencyListError(
                    ('Invalid valency for atom {symbol} ({type}) with {radicals} unpaired electrons, '
                    '{lonePairs} pairs of electrons, {charge} charge, and bonds [{bonds}].'
                    ).format(symbol=atom.symbol,
                             type=getAtomType(atom, atom.edges).label,
                             radicals=atom.radicalElectrons,
                             lonePairs=atom.lonePairs,
                             charge=atom.charge,
                             bonds=','.join([str(bond.order) for bond in atom.bonds.values()])
                            ))

    @staticmethod
    def check_multiplicity(nRad, multiplicity):
        '''
        Check that the multiplicity complies with the formula: m = 2s + 1,
        where s is the sum of the spin [+/- (1/2) ] of the unpaired electrons
        
        For a simple radical (nRad = 1): 
        s = +1/2 , m = 2 (doublet)
        
        For a biradical, s can be either 0 [+0.5 + (-0.5) ] or 1 [+0.5 + (+0.5) ] 
        and m = 1 (singlet) or m = 3 (triplet).
        '''
        if nRad in [0,1]:
            if multiplicity != (nRad + 1):
                raise InvalidAdjacencyListError('Multiplicity {0} not in agreement with total number of radicals {1}.'.format(multiplicity, nRad))
        elif nRad == 2:
            if not int(multiplicity) in [1,3]: raise InvalidAdjacencyListError('Multiplicity {0} not in agreement with total number of radicals {1}.'.format(multiplicity, nRad))
        elif nRad == 3:
            if not int(multiplicity) in [4,2]: raise InvalidAdjacencyListError('Multiplicity {0} not in agreement with total number of radicals {1}.'.format(multiplicity, nRad))
        elif nRad == 4:
            if not int(multiplicity) in [5,3,1]: raise InvalidAdjacencyListError('Multiplicity {0} not in agreement with total number of radicals {1}.'.format(multiplicity, nRad))
        else: logging.warning("Consistency checking of multiplicity of molecules with more than 4 unpaired electrons is not implemented yet!")
    
    @staticmethod
    def check_hund_rule(atom, multiplicity):
        '''
        It is checked whether atoms with 2 unpaired electrons on the same atom
        result in a multiplicity of 3, and not 1. 
        
        Unpaired electrons in 2 different orbitals belonging to the same atom
        should have the same spin, and hence, should result in a multiplicity of 3. 
        '''
        if atom.radicalElectrons == 2 and multiplicity == 1:
            raise InvalidAdjacencyListError("Violation of hund's rule. Invalid multiplicity of {0} because there is an atom with {1} unpaired electrons"
                                            .format(multiplicity, atom.radicalElectrons))
            
################################################################################

def fromOldAdjacencyList(adjlist, group=False, saturateH=False):
    """
    Convert a pre-June-2014 string adjacency list `adjlist` into a set of :class:`Atom` and
    :class:`Bond` objects. 
    It can read both "old style" that existed for years, an the "intermediate style" that
    existed for a few months in 2014, with the extra column of integers for lone pairs.
    """
    atoms = []
    atomdict = {}
    bonds = {}
    
    try:
        adjlist = adjlist.strip()
        lines = adjlist.splitlines()
        if adjlist == '' or len(lines) == 0:
            raise InvalidAdjacencyListError('Empty adjacency list.')

        # Skip the first line if it contains a label
        if len(lines[0].split()) == 1:
            label = lines.pop(0)
            if len(lines) == 0:
                raise InvalidAdjacencyListError("""Error in adjacency list\n{0}\nNo atoms specified.""".format(adjlist))
        
        mistake1 = re.compile('\{[^}]*\s+[^}]*\}')
        atomicMultiplicities = {} # these are no longer stored on atoms, so we make a separate dictionary
        # Iterate over the remaining lines, generating Atom or GroupAtom objects
        for line in lines:

            # Sometimes people put spaces after commas, which messes up the
            # parse-by-whitespace. Examples include '{Cd, Ct}'.
            if mistake1.search(line):
                raise InvalidAdjacencyListError(
                    "Error in adjacency list: \n{1}\nspecies shouldn't have spaces inside braces: {0}".format(mistake1.search(line).group(), adjlist)
                    )

            # Sometimes commas are used to delimit bonds in the bond list,
            # so replace them just in case
            line = line.replace('},{', '} {')
            
            data = line.split()

            # Skip if blank line
            if len(data) == 0: continue

            # First item is index for atom
            # Sometimes these have a trailing period (as if in a numbered list),
            # so remove it just in case
            aid = int(data[0].strip('.'))

            # If second item starts with '*', then atom is labeled
            label = ''; index = 1
            if data[1][0] == '*':
                label = data[1]
                index += 1

            # Next is the element or functional group element
            # A list can be specified with the {,} syntax
            atomType = data[index]
            if atomType[0] == '{':
                atomType = atomType[1:-1].split(',')
            else:
                atomType = [atomType]
            index += 1
            
            # Next is the electron state
            radicalElectrons = []; 
            additionalLonePairs = []
            elecState = data[index].upper()
            if elecState[0] == '{':
                elecState = elecState[1:-1].split(',')
            else:
                elecState = [elecState]
            if len(elecState) == 0:
                raise InvalidAdjacencyListError("Error in adjacency list:\n{0}\nThere must be some electronic state defined for an old adjlist".format(adjlist))
            for e in elecState:
                if e == '0':
                    radicalElectrons.append(0); additionalLonePairs.append(0)
                elif e == '1':
                    radicalElectrons.append(1); additionalLonePairs.append(0)
                elif e == '2': 
                    if not group:
                        raise InvalidAdjacencyListError("Error in adjacency list:\n{0}\nNumber of radical electrons = 2 is not specific enough.  Please use 2S or 2T.".format(adjlist))
                    # includes 2S and 2T
                    radicalElectrons.append(0); additionalLonePairs.append(1)
                    radicalElectrons.append(2); additionalLonePairs.append(0)
                elif e == '2S':
                    radicalElectrons.append(0); additionalLonePairs.append(1)
                elif e == '2T':
                    radicalElectrons.append(2); additionalLonePairs.append(0)
                elif e == '3':
                    if not group:
                        raise InvalidAdjacencyListError("Error in adjacency list:\n{0}\nNumber of radical electrons = 3 is not specific enough.  Please use 3D or 3Q.".format(adjlist))
                    # includes 3D and 3Q
                    radicalElectrons.append(1); additionalLonePairs.append(1)
                    radicalElectrons.append(3); additionalLonePairs.append(0)
                elif e == '3D':
                    radicalElectrons.append(1); additionalLonePairs.append(1)
                elif e == '3Q':
                    radicalElectrons.append(3); additionalLonePairs.append(0)
                elif e == '4':
                    if not group:
                        raise InvalidAdjacencyListError("Error in adjacency list:\n{0}\nNumber of radical electrons = 4 is not specific enough. Please use 4S, 4T, or 4V.".format(adjlist))
                    # includes 4S, 4T, and 4V
                    radicalElectrons.append(0); additionalLonePairs.append(2)
                    radicalElectrons.append(2); additionalLonePairs.append(1)
                    radicalElectrons.append(4); additionalLonePairs.append(0)
                elif e == '4S':
                    radicalElectrons.append(0); additionalLonePairs.append(2)
                elif e == '4T':
                    radicalElectrons.append(2); additionalLonePairs.append(1)
                elif e == '4V':
                    radicalElectrons.append(4); additionalLonePairs.append(0)
                elif e == 'X':
                    if not group:
                        raise InvalidAdjacencyListError("Error in adjacency list:\n{0}\nNumber of radical electrons = X is not specific enough.  Wildcards should only be used for groups.".format(adjlist))
                    radicalElectrons = []
            index += 1
            
            # Next number defines the number of lone electron pairs (if provided)
            lonePairsOfElectrons = None
            if len(data) > index: