# coding:iso-8859-1

# Klagomål kan skickas till författaren Jonas Sjöbergh (jsh@nada.kth.se)
# eller så kan man vifta med handen och klaga muntligt.

# ---------- Uppgift 1 -----------
    
# det här är gissningarna

persons = ["Robert"]  # alla ord listade här kommer annoteras med "person"
places = ["Göteborg"]  # alla ord listade här kommer annoteras med "place"
companies = ["TT","LRF","Icann","Celtic"] # alla ord listade här kommer annoteras med "company"

# övriga ord kommer annoteras med "-" om de inte är namn eller "entity" om de är namn

possible_tags = ["entity", "-", "person", "place", "company"] # en lista på alla annoteringar som får förekomma


# ---------- Uppgift 2 -----------

# --- Regelmönster ---
# ett regelmönster är en lista av tuppler,
# där varje par är dels kolumnen, och dels offset från
# nuvarande position. [(0, -1)] betyder t.ex. att i kolumn 0 ska man titta
# på ordet före det som vi nu är intresserade av
# kolumnerna är:
# 0 -> ord
# 1 -> part-of-speech
# 2 -> gissningen hittills (person, entity, place etc.)
patterns = [
    #[(0,0)], # ordet självt, ger i princip listorna ovan tillbaka som regler...
    [(0,-1)], # ordet före
    [(0,1)],  # ordet efter
    #[(0,1),(0,-1),(2,0)],# orden före och efter, samt vad vi tidigare gissat för detta ord
    #[(0,2)], # ordet 2 steg efter
    #[(1,-1),(1,1)], # PoS-taggarna före och efter
    #[(2, -1)] # vilken sorts entity, före
    ]

# vilka är de största och minsta värdena för regler
# som tittar på kringliggande gissningar?
# behövs för att se vilka regler vars resultat kan ha ändrats
min_guess_offset = 0
max_guess_offset = 0
for p in patterns:
    for c in p:
        if c[0] == 2:
            if c[1] < min_guess_offset:
                min_guess_offset = c[1]
            if c[1] > max_guess_offset:
                max_guess_offset = c[1]


from sys import argv # vi vill läsa kommandoradsargument

COLS = 4 # antal kolumner i data, inklusive rätt svar
# Ex: ord + part-of-speech + gissning -> 4, med den extra kolumnen för rätt gissning
# gissningskolumnen måste vara näst sist, rätt svar sist

pm = []    # fulhack
entity = 0 # fulhack

# en radda fula globala variabler
lex = []
ids = []
invlex = []
corpus = []
for i in range(COLS):
    lex.append({}) # mappa strängar till heltal, snabbare att jämföra
    invlex.append([]) # mappa åt andra hållet
    corpus.append([]) # alla data lagras columnvis här
    ids.append(0) # några fula räknare

allrules = []  # här lagras alla möjliga regler
goodrules = [] # här lagras de regler som kommer bli svaret

SIZE = 0 # antal ord i korpus

# Läser in korpus från fil. förväntar sig lika många kolumner på varje rad
def read_corpus(filename):
    f = open(filename)
    for r in f.readlines():
        tokens = r.split()  # separera tokens med whitespace
        if len(tokens) > 0:
            for i in range(len(tokens)): # för varje kolumn
                t = tokens[i]
                if not lex[i].has_key(t): # om inte detta ord fått ett heltalsvärde
                    lex[i][t] = ids[i]    # ge det ett nytt värde
                    ids[i] += 1
                    invlex[i].append(t)   # kom ihåg vilket ord som fick detta värde
                corpus[i].append(int(lex[i][t])) # stoppa in heltalet som nästa ord i korpus, i rätt kolumn
    f.close()

    global SIZE
    SIZE = len(corpus[0]) # lagra antal ord i en global variabel (fult)

    global pm # specialhack eftersom vi bara vill titta på Named Entities, sparar massa tid
    try:
        pm.append(lex[1]["pm.nom"])
    except KeyError:
        pass
    try:
        pm.append(lex[1]["pm.gen"])
    except KeyError:
        pass

# Genererar alla möjliga regler som skulle kunna göra nytta på denna position
def generate_rules_at_pos(pos, rulelex):
    if corpus[1][pos] in pm and corpus[-1][pos] != entity: # generera inte regler som byter till entity
        # detta villkor är bara för att skippa alla icke-entities så fort vi kan, borde generellt inte vara här

        newvalue = corpus[-1][pos]  # vi vill att det ska bli rätt svar
        for p in patterns:    # pröva varje regelmall
            ok = 1
            rule = [newvalue]  # regeln ska byta till 'newvalue'
            for check in p:  # för varje villkor i regelmallen
                if 0 <= pos + check[1] < SIZE: # kolla att offset fortfarande ligger inne i korpus
                    # i kolumn check[0] på offset check[1] ska det finnas värdet corpus[check[0]][pos+check[1]]
                    rule.append((check[0], check[1], corpus[check[0]][pos+check[1]]))
                else:
                    ok = 0
                    break
            if ok:
                rulelex[pack(rule)] = 0 # hasha regeln, vi vill inte ha dubbletter

# gör en sträng av en regel, eftersom listor (som reglerna egentligen är)
# inte kan hashas i python (fulhack)
def pack(rule):
    s = str(rule[0])
    for i in range(1,len(rule)):
        for e in rule[i]:
            s = "%s %s" % (s, str(e))
    return s
# gör tvärt om mot ovanstående
def unpack(rulestr):
    s = rulestr.split()
    rule = [int(s[0])]
    i = 1
    while i < len(s):
        c = (int(s[i]), int(s[i+1]), int(s[i+2]))
        rule.append(c)
        i += 3
    return rule

# träningssteget i TBL, välj ut de regler som är bra
def learn_rules(threshold):
    global allrules
    bestscore = threshold

    # generera alla möjliga regler från de mallar vi har
    print "generera regler"
    rules = {}
    for i in range(SIZE): # för varje ord i korpus
        generate_rules_at_pos(i, rules) # kolla vilka regler som vore bra där

    N = 0
    for r in rules.keys():
        allrules.append([0, unpack(r)]) # efter att ha rensat ut dubbletter, lagra regler i en lista
        N += 1
    rules = 0 # kasta hashtabellen

    if N <= 0:
        print "inga regler kunde genereras (ointressant data)"
        return

    # räkna på hur många ställen varje regel rättar respektive inför ett fel
    # (normalt sett det mest tidskrävande steget i programmet)
    print "bedöm hur bra reglerna är,", N, "regler (detta tar låååång tid)"
    for i in range(N):
        allrules[i][0] = rule_score(allrules[i][1]) # kolla hur bra denna regel är

        # skriv ut lite info, så använder märker att det händer något
        print "score:", allrules[i][0],
        print_rule(allrules[i][1])
        
    # plocka bästa regeln, applicera den, räkna om på nytt hur bra alla
    # regler är efter ändringarna
    print "kör TBL"
    while bestscore >= threshold: # om ingen regel är bra, ge upp
        bestscore = allrules[0][0]
        best = 0
        for i in range(1, len(allrules)): # hitta bästa regeln
            if allrules[i][0] > bestscore:
                best = i
                bestscore = allrules[i][0]
        
        if allrules[best][0] >= threshold: # om den är tillräckligt bra, använd den
            print "best score:", allrules[best][0],
            print_rule(allrules[best][1])
            
            goodrules.append(allrules[best][1])  # spara regler i ordning (de måste appliceras i rätt ordning)

            # ta reda på var denna regel ändrar i datat
            positions = where(allrules[best][1])
            
            # uppdatera reglernas värden, dessa kan ha ändrats
            # (här kan man optimera så det blir betydligt snabbare)
            for p in positions:
                # dra av score före ändring
                for offset in range(-max_guess_offset, -min_guess_offset +1):
                    pos = p + offset
                    if 0 <= pos < SIZE and corpus[-1][pos] != entity and corpus[1][pos] in pm:
                        for r in range(0,len(allrules)):
                            allrules[r][0] -= rule_score_pos(allrules[r][1], p+offset)

                # ändra
                corpus[-2][p] = allrules[best][1][0]

                # lägg till score efter ändring
                for offset in range(-max_guess_offset, -min_guess_offset +1):
                    pos = p + offset
                    if 0 <= pos < SIZE and corpus[-1][pos] != entity and corpus[1][pos] in pm:
                        for r in range(0,len(allrules)):
                            allrules[r][0] += rule_score_pos(allrules[r][1], p+offset)

# kolla om en viss position uppfyller alla villkor i en regel
def applicable(rule, pos):
    if corpus[1][pos] in pm:
        # detta villkor är bara för att skippa alla icke-entities så fort vi kan, borde generellt inte vara här

        for j in range(1,len(rule)): # för varje villkor i regeln
            if 0 <= pos + rule[j][1] < SIZE: # offseten får inte ta oss utanför korpus
                if corpus[rule[j][0]][pos + rule[j][1]] != rule[j][2]:
                    # om det står fel värde på positionen med denna offset...
                    return 0
            else:
                return 0
        return 1

    else:
        # detta villkor är bara för att skippa alla icke-entities så fort vi kan, borde inte vara här
        return 0

# applicera en regel, dvs inför de ändringar som matchar regeln
def apply(rule):
    newvalue = rule[0]
    for i in range(0,SIZE):
        if newvalue != corpus[-2][i] and applicable(rule,i): # ignorera ställen där ingen ändring märks
            corpus[-2][i] = newvalue

# kolla vilka positioner som skulle kunna ändras av en regel
def where(rule):
    newvalue = rule[0]
    applied_positions = []

    for i in range(0,SIZE):
        if newvalue != corpus[-2][i] and applicable(rule,i): # om ingen ändring sker, ignorera
            applied_positions.append(i) # håll reda på var vi vill ändra
    return applied_positions

# hur många fel rättar en regel? hur många nya fel inför den?
# denna funktion returnerar hur många färre fel det blir med
# denna regel (om man ignorerar taggen "entity")
def rule_score(rule):
    score = 0
    for i in range(0,SIZE):
        if corpus[-1][i] != entity: # räkna inte entity-taggar som rätt eller fel
            if applicable(rule, i):
                if corpus[-2][i] == corpus[-1][i]:
                    if corpus[-1][i] != rule[0]:
                        score -= 1 # om det är rätt svar och vi ändrar
                elif rule[0] == corpus[-1][i]:
                    score += 1 # det är inte rätt svar och vi ändrar till rätt
    return score

# som ovan, fast bara på en viss position
def rule_score_pos(rule, pos):
    score = 0
    if corpus[-1][pos] != entity: # räkna inte entity-taggar som rätt eller fel
        if applicable(rule, pos):
            if corpus[-2][pos] == corpus[-1][pos]:
                if corpus[-1][pos] != rule[0]:
                    score -= 1 # om det är rätt svar och vi ändrar
            elif rule[0] == corpus[-1][pos]:
                score += 1 # det är inte rätt svar och vi ändrar till rätt
    return score

# kör TBL i träningsläge
def trainmain():
    if len(argv) < 3:
        print "%s <train|use|seed> <corpus> [threshold]" % argv[0]
        return
    read_corpus(argv[2])

    global entity
    entity = int(lex[-1]["entity"]) # denna rad gör så att or med entity-tagg i facit inte räknas
    # vare sig som rätt eller fel, vilket ger många regler som kanske inte alltid gör rätt...

    # entity = -1 # detta ger mer konservativa regler, dvs räknar entity-taggen som vilken tagg som helst
    
    if len(argv) > 3:
        threshold = int(argv[3])
    else:
        threshold = 1
    
    learn_rules(threshold)

    # dumpa resultatet i en fil
    f = open("rules","w")
    for r in goodrules:
        f.write("%s\n" % filepack(r))
    f.close()

# dumpa en regel i en fil. det går inte att lagra
# heltalen i filen, eftersom samma ord kanske inte har
# samma heltal när reglerna läses in igen (om man inte
# lagrar mappningen mellan ord och tal, förstås)
def filepack(rule):
    s = invlex[-1][rule[0]]
    i = 1
    for i in range(1,len(rule)):
        s = "%s %s %s %s" % (s, rule[i][0], rule[i][1], invlex[rule[i][0]][rule[i][2]])
    return s
# motsatsen till funktionen ovan
def fileunpack(rulestr):
    s = rulestr.split()
    rule = [lex[-1][s[0]]]
    i = 1
    while i < len(s):
        col = int(s[i])
        offs = int(s[i+1])
        if lex[col].has_key(s[i+2]):
            c = (col, offs, lex[col][s[i+2]])
        else:
            c = (col, offs, -1) # unknown word, denna regel borde man egentligen kasta, den kommer aldrig matcha
            # print "unknown word",s[i+2]
        rule.append(c)
        i += 3
    return rule

# läs in regler från en fil
def read_rules(filename):
    # OBS! måste ske EFTER read_corpus
    f = open(filename)
    rules = []
    for r in f.readlines():
        rules.append(fileunpack(r))
    return rules

# använd regler på oannoterat data för att märka upp detta
def usemain():
    if len(argv) < 3:
        print "%s <train|use|seed> <rules> <text>" % argv[0]
        return
    read_corpus(argv[3])
    rules = read_rules(argv[2])
    initial_guess()

    # applicera reglerna i tur och ordning
    for r in rules:
        apply(r)
    print_corpus()

# skriv ut vårt data med införda ändringar
def print_corpus():
    for i in range(0,SIZE):
        line = ""
        for c in range(COLS):
            if i < len(corpus[c]):
                line = "%s%s\t" % (line, invlex[c][corpus[c][i]])
        print line

# skriv ut en regel på skärmen
def print_rule(rule):
    cols = ["ord", "PoS", "gissning"]
    print "byt gissning till",invlex[-1][rule[0]],
    for i in range(1,len(rule)):
        print "OM:",cols[rule[i][0]],rule[i][1],"är",invlex[rule[i][0]][rule[i][2]],
    print ""

# sätt initialtillståndet, ge alla Proper Names
# taggen "entity", alla andra ord taggen "-"
def initial_guess():
    for i in range(SIZE):
        if corpus[1][i] in pm:
            corpus[-2].append(int(lex[-2]["entity"]))
        else:
            corpus[-2].append(int(lex[-2]["-"]))

# metod för att latmansannotera lite data att träna på
def seedmain():
    # istället för att annotera en korpus gissar vi ett
    # par stycken ur varje kategori och hoppas att det räcker
    if len(argv) < 3:
        print "%s <train|use|seed> <corpus>" % argv[0]
        return

    read_corpus(argv[2])
    initial_guess()

    for i in range(0,SIZE):
        # om ett ord finns i någon av listorna på saker att annotera, ge det lämplig annotering
        # Här kan man stoppa in fler sorters entities att annotera om man vill.
        # Då får man dock stoppa in de nya taggarna i listan "possible_tags"
        # också. Den hittar man under "Uppgift 1" långt upp i filen
        if invlex[0][corpus[0][i]] in persons:
            corpus[-1].append(int(lex[-1]["person"]))
        elif invlex[0][corpus[0][i]] in places:
            corpus[-1].append(int(lex[-1]["place"]))
        elif invlex[0][corpus[0][i]] in companies:
            corpus[-1].append(int(lex[-1]["company"]))
        elif corpus[1][i] in pm:
            corpus[-1].append(int(lex[-1]["entity"]))
        else:
            corpus[-1].append(int(lex[-1]["-"]))
    print_corpus()

# mata ut statistik från jämförelse med facit
def evalmain():
    if len(argv) < 4:
        print "%s eval facit testfil" % argv[0]
        return

    f1 = open(argv[2])
    f2 = open(argv[3])

    more = 1
    correct = 0
    error = 0
    matrix = {}
    while more:
        l1 = f1.readline()
        l2 = f2.readline()
        
        t1 = l1.split()
        t2 = l2.split()
        if len(t1)*len(t2) > 0:
            if t1[-1] != "-":
                if t1[-1] == t2[-1]:
                    correct += 1
                else:
                    error += 1
                if not matrix.has_key(t1[-1]):
                    matrix[t1[-1]] = {}
                if not matrix[t1[-1]].has_key(t2[-1]):
                    matrix[t1[-1]][t2[-1]] = 0
                matrix[t1[-1]][t2[-1]] += 1
        else:
            more = 0
                
    for k in matrix.keys():
        sum = 0
        cor = 0
        print "Då rätt svar var",k,"gissades:",
        for kk in matrix[k].keys():
            print kk, matrix[k][kk],";",
            sum += matrix[k][kk]
            if k == kk:
                cor = matrix[k][kk]
        print 100*cor / float(sum),"%"
    print correct,"av",correct+error," blev rätt,",100*correct/float(correct+error),"%"

def main():
    # vi vill alltid ha vårt lexikon med möjliga annoteringar redo, bygg det nu
    i = 0
    for t in possible_tags:
        lex[-1][t] = i
        invlex[-1].append(t)
        lex[-2][t] = i
        invlex[-2].append(t)
        i += 1

    if len(argv) < 2:
        print "%s <train|use|seed|eval> ..." % argv[0]
    elif argv[1] == "train":
        trainmain()
    elif argv[1] == "use":
        usemain()
    elif argv[1] == "seed":
        seedmain()
    elif argv[1] == "eval":
        evalmain()

main()