Reddit Intermediate 90

#!/usr/bin/python3
# -*- coding: utf-8 -*-

unit_names = ('hertz', 'newton', 'pascal', 'joule', 'watt', 'coulomb',
              'volt', 'farady', 'ohm', 'siemens', 'weber', 'tesla', 'henry',
              'lumen', 'lux', 'becquerel', 'gray', 'sievert', 'katal', 'mass', 'velocity', 'celsius')
unit_types = ('s^-1', 'kg m s^-2', 'kg m^-1 s^-2', 'kg m^2 s^-2', 'kg m^2 s^-3', 's A',
              'kg m^2 s^-3 A^-1', 'kg^-1 m^-2 s^4 A^2', 'kg m^2 s^-3 A^-2', 'kg^-1 m^-2 s^3 A^2',
              'kg m^2 s^-2 A^-1', 'kg s^-2 A^-1', 'kg m^2 s^-2 A^-2', 'cd', 'm^-2 cd', 's^-1',
              'm^2 s^-2', 'm^2 s^-2', 's^-1 mol', 'kg', 'm s^-1', 'K')

def make_keys():
    assert len(unit_names) == len(unit_types)
    sort = lambda u: ' '.join(sorted(u.split()))
    name_to_unit = {k: v for k, v in zip(unit_names, unit_types)}
    unit_to_name = {k: v for k, v in zip(map(sort, unit_types), unit_names)}
    return name_to_unit, unit_to_name

name_to_unit, unit_to_name = make_keys()

def count_segment(seg):
    tmp = {}
    if '*' in seg:
        seg = seg.split('*')
    else:
        seg = seg.split()
    for i in seg:
        try:
            b, count = i.split('^')
        except ValueError:
            b = i.split('^')[0]
            count = None
        try:
            t = tmp[b]
        except KeyError:
            t = 0
        if count:
            tmp[b] = int(count) + t
        else:
            tmp[b] = 1 + t
    return tmp

def check_name(*args):
    new = []
    for a in args:
        try:
            new.append(name_to_unit[a])
        except KeyError:
            new.append(a)
    return ' '.join(new)

def determine_count(i, t):
    if t == 1:
        return i
    elif t:
        return '%s^%d' % (i, t)
    return None

def add_others(top, bot):
    t = set(top.keys())
    b = set(bot.keys())
    new = []
    for k in t - b:
        tmp = top[k]
        new.append(determine_count(k, tmp))
    return new

def check_units(s):
    tmp = ' '.join(sorted(s.split()))
    if tmp in unit_to_name:
        return s + ' (%s)' % (unit_to_name[tmp])
    return s

def translate(exp):
    if '/' in exp:
        top, bottom = exp.split('/')
        top, bottom = check_name(*top.split('*')), check_name(*bottom.split('*'))
        top_c, bot_c = count_segment(top), count_segment(bottom)
        both = set(top_c.keys()) & set(bot_c.keys())
        new = []
        for b in both:
            tmp = top_c[b] - bot_c[b]
            new.append(determine_count(b, tmp))
        new.extend(add_others(top_c, bot_c))
        new.extend(add_others(bot_c, top_c))

    else:
        exp = check_name(*exp.split('*'))
        count = count_segment(exp)
        new = []
        for i in count:
            tmp = count[i]
            new.append(determine_count(i, tmp))

    return check_units(' '.join(a for a in new if a))