'''Exploring the claim that a small dictionary can seed
an otherwise unsupervised HMM to learn a decent POS-tagger'''
import nltk, random
from nltk.corpus import brown
from nltk.tag.hmm import HiddenMarkovModelTagger, HiddenMarkovModelTrainer
from nltk.probability import FreqDist,ConditionalFreqDist
from nltk.probability import MLEProbDist, RandomProbDist, DictionaryConditionalProbDist

trainTagsPercent=0.99
trainHMMPercent=0.9
knownWordsPercent=0.99

TAGSETS={
  'univ':
  [u'ADJ', u'ADP', u'ADV', u'CONJ', u'DET', u'NOUN', u'NUM',
   u'PRON', u'PRT', u'VERB', u'X', u'.'],
  'brown':
  [u"ABL", u"ABN", u"ABX", u"AP", u"AP$", u"AP+AP", u"AT", u"BE",
   u"BED", u"BED*", u"BEDZ", u"BEDZ*", u"BEG", u"BEM", u"BEM*",
   u"BEN", u"BER", u"BER*", u"BEZ", u"BEZ*", u"CC", u"CD",
   u"CD$", u"CS", u"DO", u"DO*", u"DO+PPSS", u"DOD", u"DOD*",
   u"DOZ", u"DOZ*", u"DT", u"DT$", u"DT+BEZ", u"DT+MD", u"DTI",
   u"DTS", u"DTS+BEZ", u"DTX", u"EX", u"EX+BEZ", u"EX+HVD", u"EX+HVZ",
   u"EX+MD", u"FW-*", u"FW-AT", u"FW-AT+NN", u"FW-AT+NP", u"FW-BE", u"FW-BER",
   u"FW-BEZ", u"FW-CC", u"FW-CD", u"FW-CS", u"FW-DT", u"FW-DT+BEZ", u"FW-DTS",
   u"FW-HV", u"FW-IN", u"FW-IN+AT", u"FW-IN+NN", u"FW-IN+NP", u"FW-JJ",
   u"FW-JJR", u"FW-JJT", u"FW-NN", u"FW-NN$", u"FW-NNS", u"FW-NP", u"FW-NPS",
   u"FW-NR", u"FW-OD", u"FW-PN", u"FW-PP$", u"FW-PPL", u"FW-PPL+VBZ",
   u"FW-PPO",  u"FW-PPO+IN", u"FW-PPS", u"FW-PPSS", u"FW-PPSS+HV", u"FW-QL",
   u"FW-RB",  u"FW-RB+CC", u"FW-TO+VB", u"FW-UH", u"FW-VB", u"FW-VBD",
   u"FW-VBG",  u"FW-VBN", u"FW-VBZ", u"FW-WDT", u"FW-WPO", u"FW-WPS", u"HV",
   u"HV*",  u"HV+TO", u"HVD", u"HVD*", u"HVG", u"HVN", u"HVZ", u"HVZ*", u"IN",
   u"IN+IN",  u"IN+PPO", u"JJ", u"JJ$", u"JJ+JJ", u"JJR", u"JJR+CS", u"JJS",
   u"JJT",  u"MD", u"MD*", u"MD+HV", u"MD+PPSS", u"MD+TO", u"NN", u"NN$",
   u"NN+BEZ", u"NN+HVD", u"NN+HVZ", u"NN+IN", u"NN+MD", u"NN+NN", u"NNS",
   u"NNS$", u"NNS+MD", u"NP", u"NP$", u"NP+BEZ", u"NP+HVZ", u"NP+MD",
   u"NPS", u"NPS$", u"NR", u"NR$", u"NR+MD", u"NRS", u"OD",
   u"PN", u"PN$", u"PN+BEZ", u"PN+HVD", u"PN+HVZ", u"PN+MD", u"PP$",
   u"PP$$", u"PPL", u"PPLS", u"PPO", u"PPS", u"PPS+BEZ", u"PPS+HVD",
   u"PPS+HVZ", u"PPS+MD", u"PPSS", u"PPSS+BEM", u"PPSS+BER", u"PPSS+BEZ",
   u"PPSS+BEZ*", u"PPSS+HV", u"PPSS+HVD", u"PPSS+MD", u"PPSS+VB", u"QL",
   u"QLP",  u"RB", u"RB$", u"RB+BEZ", u"RB+CS", u"RBR", u"RBR+CS", u"RBT",
   u"RN",  u"RP", u"RP+IN", u"TO", u"TO+VB", u"UH", u"VB", u"VB+AT",
   u"VB+IN", u"VB+JJ", u"VB+PPO", u"VB+RP", u"VB+TO", u"VB+VB", u"VBD",
   u"VBG", u"VBG+TO", u"VBN", u"VBN+TO", u"VBZ", u"WDT", u"WDT+BER",
   u"WDT+BER+PP", u"WDT+BEZ", u"WDT+DO+PPS", u"WDT+DOD", u"WDT+HVZ", u"WP$",
   u"WPO", u"WPS", u"WPS+BEZ", u"WPS+HVD", u"WPS+HVZ", u"WPS+MD", u"WQL",
   u"WRB", u"WRB+BER", u"WRB+BEZ", u"WRB+DO", u"WRB+DOD", u"WRB+DOD*",
   u"WRB+DOZ", u"WRB+IN", u"WRB+MD",
   u"(", u")", u"*", u",", u"--", u".", u":"],
  'upenn':
  [u"CC", u"CD", u"DT", u"EX", u"FW", u"IN", u"JJ", u"JJR", u"JJS", u"LS",
   u"MD", u"NN", u"NNP", u"NNPS", u"NNS", u"PDT", u"POS", u"PRP", u"PRP$",
   u"RB", u"RBR", u"RBS", u"RP", u"SYM", u"TO", u"UH", u"VB", u"VBD", u"VBG",
   u"VBN", u"VBP", u"VBZ", u"WDT", u"WP", u"WP$", u"WRB",
   u"``", u"$", u"''", u"(", u")", u",", u"--", u".", u":"]}

TAGSETS['universal']=TAGSETS['univ']
TAGSETS['penn']=TAGSETS['upenn']

def setup(cat='news',tagset='brown',corpus=brown):
  return ([[(word.lower(),tag) for (word,tag) in s]
           for s in corpus.tagged_sents(categories=cat,tagset=tagset)],
          [(word.lower(),tag) for (word,tag) in corpus.tagged_words(categories=cat,tagset=tagset)],
          TAGSETS[tagset])

def notCurrent(s,missList):
  global i,n,done
  if done or (missList[i] is not s):
    return True
  else:
    i+=1
    if i==n:
      done=True
    return False

def splitData(words,wordPercent,sentences,sentPercent):
  global i,n, done
  trainWords=random.sample(words,int(wordPercent*len(words)))
  # random.sample(sentences,int(sentPercent*len(sentences)))
  trainSents=[s for s in sentences if random.random()<sentPercent]
  # hack!
  i=0
  n=len(trainSents)
  done=False
  testSents=[s for s in sentences if notCurrent(s,trainSents)]
  return trainWords, trainSents, testSents

def pickWords(tagged,percent):
  #wToT=ConditionalFreqDist(tagged)
  tToW=ConditionalFreqDist((t,w) for (w,t) in tagged)
  #print len(tToW[u'ADV'])
  return dict((tag,(lambda wl,p=percent:\
                wl[:int(p*len(wl))])(
                 sorted(tToW[tag].items(),key=lambda (k,v):v,reverse=True)))
          for tag in tToW.keys())

(tagged_s,tagged_w,tagset)=setup(tagset='universal')

wordTokens=FreqDist(word for word,tag in tagged_w)
print len(wordTokens), wordTokens.N()

(trainTags,trainHMM,testHMM)=splitData(tagged_w,trainTagsPercent,
                                       tagged_s,trainHMMPercent)

knownWords=pickWords(trainTags,knownWordsPercent)

class SubsetFreqDist(FreqDist):
  def __init__(self,pairs,baseset,basecount=.05):
    dict.update(self,pairs)
    self._baseset=baseset
    self._basecount=basecount
    pn=sum(n for w,n in pairs)
    self._N=pn+((len(baseset)-len(pairs))*basecount)

  def __getitem__(self,key):
    return dict.__getitem__(self,key)

  def __missing__(self,key):
    if key in self._baseset:
      return self._basecount
    else:
      return 0

  def N(self):
    return self._N

class Tag:
  def __init__(self,tag,wordsAndCounts):
    self._tag=tag
    self._wordsAndCounts=wordsAndCounts
    self._words=set(w for w,n in wordsAndCounts)
    self._nTokens=sum(n for w,n in wordsAndCounts)
    self._nTypes=len(self._words)

  def words(self):
    return self._words

  def buildPD(self,tokens):
    self._sfd=SubsetFreqDist(self._wordsAndCounts,tokens)
    self._pd=MLEProbDist(self._sfd)

  def getSFD(self):
    return self._sfd

  def getPD(self):
    return self._pd

tags=dict((tagName,Tag(tagName,wl)) for tagName,wl in knownWords.items())
kws=dict((tagName,tag.words()) for tagName,tag in tags.items())

t2=list(filter(None,
               ((lambda i:False if not i[1] else i)
                (((tagset[i],tagset[j]),
                  kws[tagset[i]].intersection(kws[tagset[j]])),)
                for i in xrange(0,len(tagset))
                for j in xrange(i+1,len(tagset)))))

for tag in tags.values():
  tag.buildPD(wordTokens)

priors = RandomProbDist(tagset)

transitions = DictionaryConditionalProbDist(
                dict((state, RandomProbDist(tagset))
                      for state in tagset))

outputs = DictionaryConditionalProbDist(
                dict((state, tags[state].getPD())
                      for state in tagset))

model = HiddenMarkovModelTagger(wordTokens, tagset,
                transitions, outputs, priors)

print model.evaluate(testHMM)

nm=HiddenMarkovModelTrainer(states=tagset,symbols=wordTokens)

# Note that contrary to naive reading of the documentation,
#  train_unsupervised expects a sequence of sequences of word/tag pairs,
#  it just ignores the tags
nnm=nm.train_unsupervised(trainHMM,model=model,max_iterations=15,updateOutputs=False)

print nnm.evaluate(testHMM)