from __future__ import division
import datetime
import os
import re
import nltk
from nltk.corpus import wordnet as wn
from nltk.corpus.reader import WordListCorpusReader
import string

PUNCTUATION = [',','.',';',':','(',')','?','!','"','\'']
CORRECTIONS = [
                ('Hie', 'the'),
                ('hie', 'the'),
                ('tho', 'the'),
                ('tbo', 'the'),
                ('foi', 'for'),
                ('fiom', 'from'),
                ('bo', 'be'),
                ('ho', 'he'),
                ('Ho', 'he'),
                ('Tho', 'The'),
                ('thc', 'the'),
                ('Thc', 'Thc'),
                ('wo', 'we'),
                ('aro', 'are'),
                ('ot', 'of'),
                ('moro', 'more'),
                ('havo', 'have'),
                ('ns', 'as'),
                ('lo', 'to'),
                ('und', 'and'),
                ('ono', 'one'),
                ('th', 'the'),
                ('Th', 'The'),
                ('ls', 'is'),
                ('cf', 'of'),
                ('ol', 'of'),
                ('tlie', 'the'),
                ('Tlie', 'The')
                ]

# Words that Voyant doesn't filter out.
# TO-DO: Need to work out why words like during and coming are being split
MY_STOPWORDS = ['Mr', 'Mrs', 'Miss', 'ing', 'ment', 'ed', 'ii', 'II', 'dur', 'said', 'com']
ENDINGS = ['ing', 'ment', 'ed', 'tion']

def clean_text(text):
    '''
    Does some basic cleaning up of punctuation and other stuff.
    '''
    #Insert spaces after full stops
    clean_text = re.sub(r'(\w)\.(.)', r'\1\. \2', text)
    # Join up hyphenated words across line breaks
    clean_text = re.sub(r'(\w)\-\s+(\w)', r'\1\2', clean_text)
    # Break up compound words
    clean_text = re.sub(r'(\w)\-(\w)', r'\1 \2', clean_text)
    # Remove odd punctuation
    clean_text = re.sub(r'\.(\w)', r' \1', clean_text)
    clean_text = re.sub(r'\-(\w)', r' \1', clean_text)
    clean_text = re.sub(r',(\w)', r' \1', clean_text)
    clean_text = re.sub(r'\'(\w)', r' \1', clean_text)
    #remove multiple spaces
    clean_text = re.sub(r'\s+', ' ', clean_text)
    # Join up some rogue endings
    for end in ENDINGS:
        clean_text = re.sub(r'(\w)\s+' + end, r'\1' + end, clean_text)
    return clean_text

def correct_text(text):
    for correction in CORRECTIONS:
        pattern = re.compile(r'\b%s\b' % correction[0])
        text = pattern.sub(correction[1], text)
    return text

def clean_downloads(path, correct=True, stopwords=False):
    '''
    Clean all text files from a Trove harvest.
    Saves cleaned versions in a new directory.
    '''
    output_dir = os.path.join(path, 'cleaned')
    report_dir = os.path.join(path, 'reports')
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    if not os.path.exists(report_dir):
        os.makedirs(report_dir)
    dirs = [ dir for dir in os.listdir(path) if os.path.isdir(os.path.join(path, dir)) and dir != 'cleaned' and dir != 'reports']
    for dir in dirs:
        print 'Processing: %s' % dir
        old_dir = os.path.join(path, dir)
        new_dir = os.path.join(output_dir, dir)
        if not os.path.exists(new_dir):
            os.makedirs(new_dir)
        report_path = os.path.join(report_dir, '%s-accuracy.txt' % dir)
        with open(report_path, 'w') as report:
            report.write('Generated: %s\n' % datetime.datetime.now().isoformat())
        entities = []
        try:
            with open(os.path.join(path, 'entities.txt'), 'r') as entities_file:
                entities = nltk.word_tokenize(entities_file.read())
        except IOError:
            try:
                with open(os.path.join(old_dir, 'entities.txt'), 'r') as entities_file:
                    entities = nltk.word_tokenize(entities_file.read())
            except IOError:
                print 'No entities file'
        files = [ os.path.join(old_dir, text_file) for text_file in os.listdir(old_dir) if text_file[-4:] == '.txt' and text_file != 'entities.txt' ]
        clean = []
        dirty = []
        total_accuracy = 0
 
       with open(report_path, 'a') as report:
            for text_file in files:
                print 'Cleaning %s' % text_file
                with open(text_file, 'r') as text:
                    content = text.read()
                    if stopwords:
                        cleaned = remove_stopwords(content)
                    cleaned = clean_file(content, entities, correct)
                accuracy = (cleaned['recognised'] / cleaned['total']) * 100
                print 'Accuracy: %s%%
' % accuracy
                total_accuracy += accuracy
                report.write('%s: %s%%\n' % (os.path.basename(text_file), accuracy))
                #clean.extend(cleaned['clean'])
                #dirty.extend(cleaned['dirty'])
                c_text = (' ').join([ word for word in cleaned['clean']])
                c_text = adjust_punctuation(c_text)
                with open(os.path.join(new_dir, os.path.basename(text_file)), 'w') as c_file:
                    c_file.write(c_text)                    
            average_accuracy = total_accuracy / len(files)
            print 'Average accuracy: %s%%' % average_accuracy
            report.write('Average accuracy: %s%%' % average_accuracy)

def adjust_punctuation(text):
    '''
    Gets rid of space between words and puctuation.
    '''
    for punc in PUNCTUATION:
        text = re.sub(r'(\w)\s+\%s' % punc, r'\1%s' % punc, text)
    return text
                

def clean_files(path):
    '''
    Clean files in the supplied directory
    '''
    texts = [text for text in os.listdir(path) if text[-4:] == '.txt']
    try:
        with open(os.path.join(path, 'entities.txt'), 'r') as entities_file:
            entities = nltk.word_tokenize(entities_file.read())
    except AttributeError:
        print 'No entities file'
        entities = []
    clean = []
    dirty = []
    for text in texts:
        print 'Cleaning %s' % text
        with open(os.path.join(path, text), 'r') as text_file:
            cleaned = clean_file(text_file.read(), entities)
            clean.extend(cleaned['clean'])
            dirty.extend(cleaned['dirty'])
    print sorted(set(dirty))
    print sorted(set(clean))
    print 'Dirty: %s' % len(dirty)
    print 'Clean: %s' % len(clean)
    
def clean_file(text, entities, correct):
    clean = []
    recognised = 0
    dirty = []
    ctext = remove_html(text)
    ctext = clean_text(ctext)
    if correct: ctext = correct_text(ctext)
    tokens = nltk.word_tokenize(ctext)
    unusual = unusual_words(tokens)
    for token in tokens:
        ltoken = token.lower().strip()
        if ltoken in PUNCTUATION:
            clean.append(token)
            recognised += 1
        else:
            if len(ltoken) > 1 or ltoken in ['a', 'i']:
                if ltoken in unusual:
                    if len(ltoken) > 3:
                        stem = wn.morphy(ltoken)
                        if stem:
                            clean.append(token)
                            recognised += 1
                        else:
                            if ltoken in entities:
                               clean.append(token)
                               recognised += 1
                            else:
                                dirty.append(token)
                                clean.append('[[%s]]' % token)
                    else:
                        dirty.append(token)
                        clean.append('[[%s]]' % token)
                else:
                    clean.append(token)
                    recognised += 1
            else:
                dirty.append(token)
    return {'clean': clean, 'dirty': dirty, 'recognised': recognised, 'total': len(tokens)}
    
def remove_html(text):
    text = text.replace('<p>','').replace('</p>','').replace('<span','').replace('</span>','')
    return text
        
def unusual_words(text):
    '''
    Returns a list of words in the supplied text than don't occur in NLTK's wordlist.
    '''
    #print text
    text_vocab = set(get_lemma(w.lower()) for w in text)
    known_words = []
    known_words.extend(word.