# Parsing of PDB chemical formulae.
#
# Author       : Chris X. Weichenberger
# Created      : 2012-07-20
# License
# Twilight is a GUI for displaying ligands incompatible with electron density.
#    Copyright (C) 2012  Christian X. Weichenberger, Edwin Pozharski, and
#    Bernhard Rupp.
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Modifications:
# 2012-07-20 (CW) Object is a list of tuples (element, nr).
# 2012-07-25 (CW) Expanded charge parsing. __str__ method.
# 2012-08-02 (CW) Computation of molecular weight.
#

import re

class CChemicalFormulaData(list):
    """Helper class to parse chemical formula from PDB.
    Example: 'C6 H13 N O4 S', or 'MN 2', or 'CL -1', or 'MG 2+' """
    def __init__(s, formula):
        s.NrHeavy    = 0;
        s.NrHydrogen = 0;
        s.Charge     = 0;
        paranthesis  = re.match(".*\((.+)\)", formula);
        if( paranthesis!=None ):
            formula = paranthesis.groups()[0]
        for compound in formula.split():
            # Check for the charge. This could be '-1' or '2+' or '3'
            mObj = re.match("([+-]*)([0-9])([+-]*)", compound);
            if( mObj!=None ):
                sig1, charge, sig2 = mObj.groups();
                # the sign is either before or after.
                sign = sig1 + sig2;
                try:
                    s.Charge = int(charge);
                    if( sign=="-" ):
                        s.Charge = -s.Charge;
                    continue ;
                except ValueError:
                    s.Charge = 0;
            mObj     = re.match("(^[A-Z]+)([0-9]*)", compound);
            elem, nr = mObj.groups()
            elem     = elem.upper();
            if( nr=="" ):
                nr = 1;
            nr = int(nr);
            s.append( (elem, nr) );
            if( elem=="H" ):
                s.NrHydrogen += nr;
            else:
                s.NrHeavy    += nr;

    def __str__(s):
        st = " ".join(["%s%d" % (elem, nr) for elem, nr in s]);
        if( s.isCharged ):
            st = st + " " + str(s.getCharge());
        return st;
    
    def isCharged(s):
        return s.Charge!=0;

    def getCharge(s):
        return s.Charge;
    
    def nrAtoms(s):
        return s.NrHydrogen + s.NrHeavy;

    def nrHeavyAtoms(s):
        return s.NrHeavy;

    def nrHydrogenAtoms(s):
        return s.NrHydrogen;


if( __name__=="__main__" ):
    line = '2(C3 H8 O)';
    #line = 'MG 2+';
    #line = 'CL -1';
    print CChemicalFormulaData(line);
    
