# Created:20050107
# By Jeff Connelly

# Balance chemical equations (perform stoichiometry)
#
# # ~/mnt/ad4s2d/jeff/slo/2005/Winter/CHEM124\ -\ Hascall/balance.py

# Copyright (C) 2005 Jeff Connelly

# This is called by a CGI script, stoich.cgi

# a*Fe(NO3)3(aq) + b*Na2CO3(aq) = c*Fe2(CO3)3(aq) + d*NaNO3(aq) 

# a*Fe = c*Fe*2		a=2*c
# a*NO3*3 = d*NO3	3*a=d
# b*Na*2 = d*Na		2*b=d
# b*CO3 = c*CO3*3	b=3*c

# solve for integers a,b,c,d

import types


# split_terms("a + b + c") => ["a", "b", "c"]
def split_terms(expression):
    terms = expression.split("+")  # Note: doesn't do subtraction
    for i in range(0, len(terms)):
        terms[i] = terms[i].strip()   # allow spaces between +
    return terms

# split_eqn("a + b = c") => ["a + b", "c"]
def split_eqn(eqn):
    if chr(63) in eqn:
        # If you paste the arrow symbol (-->, converted) from MSWord
        # into a text document and save as an ANSI-encoded file, you get this
        equalities = eqn.split(chr(63))
    elif "->" in eqn:       # poor man's arrow
        equalities = eqn.split("->")
    elif "→" in eqn:  # Unicode arrow
        equalities = eqn.split("→")
    else:
        equalities = eqn.split("=")  # Note: doesn't do inequalities
    for i in range(0, len(equalities)):
        equalities[i] = equalities[i].strip()
    return equalities

# split_state("blah(aq)") => ["blah", "aq"]
def split_state(s):
    molecule = s
    state = None
    if "(aq)" in s: state = "aq"  # aqueaous
    if "(g)" in s: state = "g"    # gaseous
    if "(l)" in s: state = "l"    # liquid
    if "(s)" in s: state = "s"    # solid
    
    if "(aq)" in s: molecule = s[:s.index("(aq)")]
    if "(g)" in s: molecule = s[:s.index("(g)")]
    if "(l)" in s: molecule = s[:s.index("(l)")]
    if "(s)" in s: molecule = s[:s.index("(s)")]

    return molecule, state


# Read an element from a molecule string
def get_element(molecule, at):
    if at >= len(molecule): return None, None

    elt = ""
    if molecule[at] == "(":
        return get_group(molecule, at)

    if not molecule[at].isupper():
        print "Tried to read element, but found nothing. ", molecule, at
        print "Found:", molecule[at]
        raise SystemExit

    # An element begins with capital, followed by lowercase
    while True:
        elt += molecule[at]
        at += 1
        if at >= len(molecule):
            break
        if not molecule[at].islower():
            break

    rep, at = get_number(molecule, at)

    print "@=",at
    return (elt, rep), at


# Read digits from the molecule string, if they are there at current position
# Always safe to call after get_element to find how many times it is there,
# also useful after a group
def get_number(molecule, at):
    if at >= len(molecule): return 1, at

    digs = ""
    while molecule[at].isdigit():
        digs += molecule[at]
        at += 1
        if at >= len(molecule):
            break
    if digs == "": 
        digs = 1   # default
    return int(digs), at

# True if composed of multiple ions
def is_polyatomic(ion_str):
    caps = 0
    for i in ion_str:
       caps += i.isupper()
    return caps > 1

# Recursively print a tree ion structure
def str_ions(ions):
    s = ""
    for i in ions:
        if type(i) == types.TupleType and \
           type(i[0]) == types.StringType and \
           type(i[1]) == types.IntType: # ("X", N)
           if i[1] == 1:
               s += i[0]
           else:
               s += "%s%d" % (i[0], i[1])
        elif type(i) == types.TupleType:
            s += "(%s)" % str_ions(i)
        elif type(i) == types.IntType:
            if i != 1: s += str(i)
        else:
            print "Unknown type: ", type(i)
            raise SystemExit
    return s

# Look for and collapse polyatomic ions
# TODO: Recursive, and don't nest so many levels deep
# XXX: Not used, instead get_element reads if polyatomic
def collapse_poly(ions):

    # (("N", 1), ("O", "2")) => ("NO2", 1)
    n = 0
    new_ions = []
    print "COLLAPSING: ", ions
    while True:
        poly = False

        # Works for the most part but current problem is that
        # "Fe(NO2)3" => ((('Fe', 1), ((('NO2', 1),), 3)), 1) 
        # Should instead be (('Fe', 1), ('NO2', '3')) -- not so round-about!
 
        # Collapse ((("NO2", 1),) -> ("NO2", 1)
        if 0 and type(ions[n][0]) == type((1,2,3)) and len(ions[n][0]) == 1:
            new_ions.append((ions[n][0][0]))
            n += 1
            if n >= len(ions): break
            continue

        if 0 and type(ions[n][0]) == type((1,2,3)): # tries to be recursive..
            new_ions.append(collapse_poly(ions[n][0]))
            n += 1
            if n >= len(ions): break
            continue

        print type(ions[n]), type(ions[n][0])
        if n + 1 < len(ions):
            if ions[n] == ("N", 1) and ions[n + 1] == ("O", 2):
                print "FOUND NO2!!!!"
                new_ions.append(("NO2", 1))
                n += 1
                poly = True
        if not poly: #  didn't match any polyatomic ion patterns
            new_ions.append(ions[n])
        n += 1
        if n >= len(ions): break
    return new_ions

def get_all(molecule, at=0):
    s = []
    last_elt = None
    while True:
        elt, at = get_group(molecule, at)
        if elt == last_elt: break # Bad way to terminate, but get_group doesn't seem to update
        s.append(elt)
        last_elt = elt
    return s

# Read a parenthetical group, like (NO3)3 -- possibly including rep
def get_group(molecule, at):
    global poly
    group = ""
    if molecule[at] != "(":
        group = molecule[at:]  # no group here
        group_rep = 1
    else:
        at += 1
        level = 1
        while level != 0 and at < len(molecule):
            if molecule[at] == "(": level += 1
            if molecule[at] == ")": level -= 1
            if level == 0: break
            group += molecule[at]
            at += 1
        if molecule[at] == ")": at += 1

        group_rep, at = get_number(molecule, at)

    # Read group into ions now
    iat = 0
    ions = []
    while iat < len(group):
        print "Getting element at ",iat
        elt, iat = get_element(group, iat)
        # Look ahead. Is the next element + this one part of a polyatomic sequence?
        #elt2, iat2 = get_element(group, iat)
        #for p in poly:
        #    if (elt, elt2) == p:
        #        print "MATCHED"
        #        elt = str_ions(p)
        #        iat = iat2  # Use this element
        #        break
        #print "2222: ", elt,elt2

        if type(elt[0]) == types.TupleType and elt[0][1] == 1:
            # ((X, 1), N)  => (X, N). The need for t his conversion is an artifact of how
            # polyatomic ions are detected
            elt = elt[0][0], elt[1]
        print "ION:", elt
        ions.append(elt)

    #ions = collapse_poly(ions)
    print str_ions(ions)

    #return (tuple(ions) + (group_rep, )), at
    return (tuple(ions), group_rep), at 


# split_m("Fe(NO3)3(aq)") => {"Fe": 1, "NO3": 3}
def split_m(m):
    molecule, state = split_state(m)
    # We ignore the state, not advanced enough to derive it, and
    # it is simple enough to do manually.

    at = 0
    while True:
        # Elements begin with a capital letter
        if not molecule[at].isupper():
            print "Got confused parsing",molecule,"at",at
            raise SystemExit

        # Maybe its a group
        molecule="Fe(NO2)3"
        elt, at = get_group(molecule, at)
        print elt, at
        raise SystemExit


#print split_state("Fe(NO3)3(aq)")

global poly
poly = []
# Polyatomic ions - don't actually need to recognize these
poly_ions = ["NH4", "H3O", "CH3OO", "C2H3O2", "CN", "OH", "ClO", "ClO2", \
             "ClO3", "ClO4", "NO2", "NO3", "MnO4", "CO3", "HCO3", "CrO4", \
             "Cr2O7", "O2", "PO4", "HPO4", "H2PO4", "SO3", "SO4", "HSO4"]
def pinit():
    global poly
    poly = [ (("N", 1), ("H", 4), 1), 
             (("H", 3), ("O", 1), 1), 
             (("C", 1), ("H", 3), ("O", 2), 1),
             (("C", 2), ("H", 3), ("O", 2), 1),  
             (("C", 1), ("N", 1), 1), 
             (("O", 1), ("H", 1), 1), 
             (("Cl",1), ("O", 1), 1), 
             (("Cl",1), ("O", 2), 1),
             (("Cl",1), ("O", 3), 1),
             (("Cl",1), ("O", 4), 1),
             (("N", 1), ("O", 2), 1),
             (("Mn",1), ("O", 4), 1),
             (("C", 1), ("O", 3), 1),
             (("H", 1), ("C", 1), ("O", 3), 1),
             (("Cr",1), ("O", 4), 1),
             (("Cr",2), ("O", 7), 1),
             (("O", 2), 1),
             (("P", 1), ("O", 4), 1),
             (("H", 1), ("P", 1), ("O", 4), 1),
             (("H", 2), ("P", 1), ("O", 4), 1),
             (("H", 2), ("P", 2), ("O", 4), 1),
             (("S", 1), ("O", 3), 1),
             (("S", 1), ("O", 4), 1),
             (("H", 1), ("S", 1), ("O", 4), 1),
             ]

pinit()

# Return a hash with counts of the elements in a term, wrapper for recursive counte
# XXX: TODO
def counte_term(term):
    counts = {}
    for molecule in term:
        # Take next, and multiply by
        rep = molecule[1]
        for element in molecule[0]:
            counts = comb(counts, counte(element))
    return counts

# Combine two hashes additively
# NOTE: Slow, for faster use sets:
#<The-Fixer> >>> if __import__('sys').version_info < (2,4): from sets import Set as set
#<The-Fixer> >>> d3 = d1.copy(); d3.update(d2)
#<The-Fixer> >>> for k in set(d1) & set(d2): d3[k] = d1[k]+d2[k]
def comb(a,b):
    for k in a:
        if k in b:
            a[k] += b[k]
            b.pop(k)
    a.update(b)
    return a

# Combine two hashes multiplicitively
def hmult(a,b):
    for k in a:
        if k in b:
            a[k] += b[k]
            b.pop(k)
        a.update(b)
    return a

#print comb({1:2,3:4},{1:3})
#raise SystemExit

# Recursively multiply and count each term
# XXX: TODO: Work!
def counte(term, mult=1):
    h = {}
    if type(term) != types.TupleType:
        print "Trying to count not tuple!", type(term), term
        raise SystemExit

    #print "COUNTING: ",term
    # TODO: Make this work

    # ("Element" or tuple, count) - repeat count times
    if type(term[1]) == types.IntType:
        #print "REPEAT:",term[0],"#",term[1]
        mult *= term[1]
        if type(term[0]) == types.StringType:
            h = {term[0]: mult}
        else:
            h = counte(term[0], mult)
    # (tuple, tuple) = no repetition, treat tuples separately. Could also be len>2
    elif type(term[0]) == types.TupleType and type(term[1]) == types.TupleType or len(term) > 2:
        #print "Multiplicity"
        for t in term:
            h = comb(h, counte(t, mult))
    return h
# We need a procedure to traverse the tree and count the elements to balance the tree!
# Call it per term, then find coefficients, then solve. Now called counte.

# Format the solution by prepending the coefficients to the terms!
def str_soln(soln, terms1s, terms2s):
    print "Solution: ", soln
    # Print solution
    vi = 0
    ps = []
    for t in terms1s:
        if soln[vi] == 1: 
            sn = ""
        else:
            sn = str(soln[vi])

        ps.append(sn + t)
        vi += 1
    fs = " + ".join(ps)
    ps = []
    for t in terms2s:
        if soln[vi] == 1:
            sn = ""
        else:
            sn = str(soln[vi])

        ps.append(sn + t)
        vi += 1
    fs += " = " + " + ".join(ps)
    return fs


def solve(e):
    eql = split_eqn(e)
    if len(eql) != 2:
        return "Not an valid chemical eqn, has "+str(len(eql))+" parts not 2"

    side1, side2 = eql

    terms1s, terms2s = split_terms(side1), split_terms(side2)
    terms1, terms2 = map(split_state, terms1s), map(split_state, terms2s)
    # Terms{1,2} are now arrays of tuples of (term, state). State needs to be
    # parsed no longer, but term needs to be parsed using get_group
    terms1 = map(lambda x: get_all(x[0]), terms1)
    terms2 = map(lambda x: get_all(x[0]), terms2)
    print 
    print

    # Each term gets a coeff var
    num_vars = len(terms1) + len(terms2)
    vars = []
    for v in range(num_vars):
        ch = chr(ord("a") + v)
        if not ch.islower():
            return "Too many vars?! Bad var: ", ch
        vars.append(ch)
    print "Using %d vars: %s" % (num_vars, str(vars))
    check = "if "

    # Get a list of all elements involved
    all_elts = {}
    for t in terms1: all_elts.update(counte_term(t))
    for t in terms2: all_elts.update(counte_term(t))

    all_elts = all_elts.keys()
    print all_elts

    for elt in all_elts:
        vi = 0
        check += "0"
        # Each term gets a alphabetic coefficient (but 0 if element not in this term)
        for t in terms1:
            ecnt = counte_term(t)
            if ecnt.has_key(elt):
                check += " + %s*%d" % (vars[vi], ecnt[elt])
            vi += 1
        check += " == 0"
        for t in terms2:
            ecnt = counte_term(t)
            if ecnt.has_key(elt):
                check += " + %s*%d" % (vars[vi], ecnt[elt])
            vi += 1
        #check +=        "\\\n   "
        check += " and "
    #check = check[:-len("\\\n   ")] + ":"
    check += "True:"
    print check

    #soln = None
    code = "def s():\n print 'Finding solution...'\n"
    vi = 1 
    # Generate code to brute-force ranges
    for v in vars:
        code += "%sfor %s in range(1,20):\n" % (" " * vi, v)
        vi += 1
    code += (" " * vi) + check + "\n"
    code += (" " * (vi + 1)) + "return " + ",".join(vars)
    print code
    eval(compile(code, "", "exec"))

    soln = eval("s()")

    if soln == None:
        return "Sorry no solution found"
   
    print
    return str_soln(soln, terms1s, terms2s)



#solve("Fe(NO3)3(aq) + NaOH(aq) = Fe(OH)3(aq) + NaNO3(aq)")
#raise SystemExit

if __name__ == "__main__":
    for a in range(1,10):
        for b in range(1,10):
            for c in range(1,10):
                for d in range(1,10):
                    if a==2*c and 3*a==d and 2*b==d and b==3*c: print a,b,c,d

    # 2Fe(NO3)3(aq) + 3Na2CO3(aq) -> Fe2(CO3)3(aq) + 6NaNO3(aq) 
    
    # SOLVED

    # Note: = is actually a ->
    # Fe(NO3)3(aq) + Na2CO3(aq) = Fe2(CO3)3(aq) + NaNO3(aq) 

    print solve("Fe(NO3)3(aq) + Na2CO3(aq) = Fe2(CO3)3(aq) + NaNO3(aq)")
    while True:
        print solve(raw_input())
    raise SystemExit

# Problem: Cu(NO3)2(aq) + Na3PO4(aq) ? Cu3(PO4)2(aq) + Na3(NO3)2  cannot be balanced??


#for i in poly_ions:
    #elt, at = get_group("(HP)2O4", 0)
    #print elt, str_ions(elt)
    #elt, at = get_group("(HP)2O4", at)
    #print elt, str_ions(elt)
    #elt, at = get_group("(HP)2O4", at)
#    elt = get_group("AgCl3", 0)[0][0]
#    print elt, str_ions(elt), elt == (("Ag", 1), ("Cl", 3))
#    raise SystemExit