#!/usr/bin/python3
"""Balance formulas, because I'm stupid and can't count to 5.

>>> Cu = elements.Cu
>>> print(Cu)
Cu
>>> print(Cu[2])
Cu₂
>>> O = elements.O
>>> print(Cu[2] + O)
Cu₂O
>>> Na, Cl, H = elements.Na, elements.Cl, elements.H
>>> pre = Cu | 2*(Na+Cl) | 2*(H[2]+O)
>>> print(pre)
Cu + 2NaCl + 2H₂O
>>> post = Cu+Cl[2] | 2*(Na+O+H) | H[2]
>>> print(post)
CuCl₂ + 2NaOH + H₂
>>> sorted(pre.totals().items())
[('Cl', 2), ('Cu', 1), ('H', 4), ('Na', 2), ('O', 2)]
>>> sorted(post.totals().items())
[('Cl', 2), ('Cu', 1), ('H', 4), ('Na', 2), ('O', 2)]
>>> imbalance(pre, post)
{}
>>> imbalance(pre | H[2], post)
{'H': 2}
>>> imbalance(pre, post | Cl[2])
{'Cl': -2}
>>> Mg + OH[2]
<Compound Mg(OH)₂>


"""

class CompoundList:
    def __or__(self, other):
        "Join compound lists."
        return Juffowup(self, other)

    def __repr__(self):
        return '<CompoundList {}>'.format(self)


class Compound(CompoundList):
    needs_parens_if_subscripted = True

    def __getitem__(self, n):
        assert isinstance(n, int)
        return Multiple(self, n)

    def __rmul__(self, n):
        "Multiplied on the left by, hopefully, an int."
        assert isinstance(n, int)
        return Proportion(self, n)

    def __repr__(self):
        return '<Compound {}>'.format(self)

    def __add__(self, other):
        "Make a bigger compound from smaller ones: Na + Cl"
        return CompoundCompound(self, other)


class Element(Compound):
    needs_parens_if_subscripted = False
    
    def __init__(self, symbol):
        self.symbol = symbol

    def __str__(self):
        return self.symbol

    def totals(self):
        return {self.symbol: 1}


class Elements:
    def __getattr__(self, symbol):
        return Element(symbol)

elements = Elements()

# Enable me to import a few elements directly:
for _element in 'Cu Na Cl H O C N Ca K Mg Al P F'.split():
    globals()[_element] = Element(_element)


class Multiple(Compound):
    "A subscripted unit like (OH)₂ or Cu₂."
    def __init__(self, monomer, n):
        assert isinstance(n, int)
        assert n > 0
        self.monomer = monomer
        self.n = n

    def __str__(self):
        unit = str(self.monomer)
        if self.monomer.needs_parens_if_subscripted:
            unit = '({})'.format(unit)

        if self.n == 1:
            return unit

        return unit + subscript_digits(self.n)

    def totals(self):
        return multiply_totals(self.n, self.monomer.totals())


_subscripts = {'0': '₀', '1': '₁', '2': '₂', '3': '₃', '4': '₄',
               '5': '₅', '6': '₆', '7': '₇', '8': '₈', '9': '₉'}

def subscript_digits(n):
    return ''.join(_subscripts[digit] for digit in str(n))


class CompoundCompound(Compound):
    "A compound of the form FooBar where Foo and Bar are compounds."
    def __init__(self, a, b):
        self.a = a
        self.b = b

    def __str__(self):
        return '{}{}'.format(self.a, self.b)

    def totals(self):
        return add_totals(self.a.totals(), self.b.totals())


class Proportion(CompoundList):
    "A compound with its stoichiometric coefficient."
    def __init__(self, compound, n):
        self.compound = compound
        assert isinstance(n, int)
        assert n > 0
        self.n = n

    def __str__(self):
        if self.n == 1:
            return str(self.compound)
        else:
            return '{}{}'.format(self.n, self.compound)

    def totals(self):
        return multiply_totals(self.n, self.compound.totals())


class Juffowup(CompoundList):
    "Two or more compounds with possible coefficients, with + between."
    def __init__(self, a, b):
        self.a = a
        self.b = b

    def __str__(self):
        return '{} + {}'.format(self.a, self.b)
    
    def totals(self):
        return add_totals(self.a.totals(), self.b.totals())


def add_totals(a, b):
    for element, amount in b.items():
        if element not in a:
            a[element] = 0
        a[element] += amount
        if a[element] == 0:
            del a[element]

    return a

def multiply_totals(n, result):
    for element in result:
        result[element] *= n
    return result


def imbalance(pre, post):
    return add_totals(pre.totals(), multiply_totals(-1, post.totals()))


H2O = H[2] + O
OH = O + H
NaOH = Na + OH
CO2 = C + O[2]
CO3 = C + O[3]
NH4 = N + H[4]
NH3 = N + H[3]
PO4 = P + O[4]
O2 = O[2]
N2 = N[2]


if __name__ == '__main__':
    import doctest
    doctest.testmod()