import string

AA = 1   ## AA   odd     AA D
AE = 2   ## AE   at      AE T
AH = 3   ## AH   hut     HH AH T
AO = 4   ## AO   ought   AO T
AW = 5   ## AW   cow     K AW
AY = 6   ## AY   hide    HH AY D
EH = 7   ## EH   Ed      EH D
ER = 8   ## ER   hurt    HH ER T
EY = 9   ## EY   ate     EY T
IH = 10  ## IH   it      IH T
IY = 11  ## IY   eat     IY T
OW = 12  ## OW   oat     OW T
OY = 13  ## OY   toy     T OY
UH = 14  ## UH   hood    HH UH D
UW = 15  ## UW   two     T UW

B  = 16  ## B    be      B IY
CH = 17  ## CH   cheese  CH IY Z
D  = 18  ## D    dee     D IY
DH = 19  ## DH   thee    DH IY
F  = 20  ## F    fee     F IY
G  = 21  ## G    green   G R IY N
HH = 22  ## HH   he      HH IY
JH = 23  ## JH   gee     JH IY
K  = 24  ## K    key     K IY
L  = 25  ## L    lee     L IY
M  = 26  ## M    me      M IY
N  = 27  ## N    knee    N IY
NG = 28  ## NG   ping    P IH NG
P  = 29  ## P    pee     P IY
R  = 30  ## R    read    R IY D
S  = 31  ## S    sea     S IY
SH = 32  ## SH   she     SH IY
T  = 33  ## T    tea     T IY
TH = 34  ## TH   theta   TH EY T AH
V  = 35  ## V    vee     V IY
W  = 36  ## W    we      W IY
Y  = 37  ## Y    yield   Y IY L D
Z  = 38  ## Z    zee     Z IY
ZH = 39  ## ZH   seizure S IY ZH ER

EMPTY_STRESS     = 0     # No stress information for that phoneme
NO_STRESS        = 0x40  # Phoneme is flagged in the dictionary as unstressed
PRIMARY_STRESS   = 0x80
SECONDARY_STRESS = 0xc0

INVALID = 0xff

code_from_string_table = {
    "AA": AA, "AE": AE, "AH": AH, "AO": AO, "AW": AW, "AY": AY,
    "EH": EH, "ER": ER, "EY": EY, "IH": IH, "IY": IY, "OW": OW,
    "OY": OY, "UH": UH, "UW": UW,
    "B" :  B, "CH": CH, "D" :  D, "DH": DH, "F" :  F, "G" :  G,
    "HH": HH, "JH": JH, "K" :  K, "L" :  L, "M" :  M, "N" :  N,
    "NG": NG, "P" :  P, "R" :  R, "S" :  S, "SH": SH, "T" :  T,
    "TH": TH, "V" :  V, "W" :  W, "Y" :  Y, "Z" :  Z, "ZH": ZH
    }

code_to_string_table = {}
for s, i in code_from_string_table.items():
    code_to_string_table[i] = s

stress_from_string_table = {
    "":  EMPTY_STRESS,
    "0": NO_STRESS,
    "1": PRIMARY_STRESS,
    "2": SECONDARY_STRESS
    }

stress_to_string_table = {}
for s, i in stress_from_string_table.items():
    stress_to_string_table[i] = s

# returns a phoneme's code and the stress separately
def split(phon):
    return phon & 0x3f, phon & 0xc0

def split_to_string(phon):
    code, stress = split(phon)
    return (code_to_string_table.get(code, "??"),
            stress_to_string_table.get(stress, "?"))

# combine a code and a stress into a phoneme
def join(code, stress):
    return (code & 0x3f) | (stress & 0xc0)

def join_from_string(code_str, stress_str):
    code = code_from_string_table.get(code_str, INVALID)
    stress = stress_from_string_table.get(stress_str, INVALID)
    if code == INVALID or stress == INVALID:
        return INVALID
    return join(code, stress)

def to_string(phon):
    code, stress = split(phon)
    return "%s%s" % (code_to_string_table.get(code, "??"),
                     stress_to_string_table.get(stress, "?"))

def from_string(s):
    if not s:
        return INVALID
    stress = stress_from_string_table.get(s[-1], EMPTY_STRESS)
    if stress != EMPTY_STRESS:
        s = s[:-1]
    code = code_from_string_table.get(s, 0)
    if not code:
        return INVALID
    return join(code, stress)

def is_vowel(phon):
    return (phon & 0x3f) <= 15

def is_consonant(phon):
    return (phon & 0x3f) >= 16

# no stressed (i.e. a 1 or a 2)
def is_stressed(phon):
    return phon & 0x80

# is flagged with any stress at all, including a 0
def is_explicitly_stressed(phon):
    return phon & 0xc0

def decomp_to_string(decomp):
    return string.join(map(to_string, decomp))

def decomp_from_string(s):
    return map(from_string, s.split())

def equal_mod_stress(phon1, phon2):
    return (phon1 & 0x3f) == (phon2 & 0x3f)

def equal_mod_slant(phon1, phon2):
    if is_vowel(phon1) and is_vowel(phon2):
        return 1
    return equal_mod_stress(phon1, phon2)