#!/usr/bin/pypy3
#coding=utf-8

#===============================================================================
#   Web Corpus Wordlist Based Language Filter
#   Copyright © 2016-2020 Vít Suchomel
#   http://corpus.tools/wiki/languagefilter
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program. If not, see <http://www.gnu.org/licenses/>.
#===============================================================================

USAGE = '''
Separate documents and paragraphs by language using word frequency lists.
All languages to recognise have to be specified and respective frequency wordlists supplied.

Usage:
./lang_filter.py (LANGUAGE FRQ_WORDLIST[.gz|.xz])+ ACCEPTED_LANGS REJECTED_OUT LANG_RATIO_THR

E.g. to see all score information in a single output file:
./lang_filter.py Czech wl/czech.frqwl.gz Slovak wl/slovak.frqwl.gz \\
    English wl/english.frqwl.gz Czech,Slovak out_rejected.vert 1.01 < in.vert > out.vert \\

Or just to output to separate documents:
./lang_filter.py Czech wl/czech.frqwl.gz Slovak wl/slovak.frqwl.gz \\
    English wl/english.frqwl.gz Czech,Slovak out_rejected.vert 1.01 < in.vert > out.vert \\
    | grep -v '^<par_langs' | cut -f1 \
    | ./vertsplit_by_attr.py doc lang out.vert.lang_

LANGUAGE language name

FRQ_WORDLIST corpus frequency wordlist
    format: <word>TAB<count>, one record per line, can be gzipped or xzipped

ACCEPTED_LANGS comma separated list of accepted languages (ALL to accept all)

REJECTED_OUT a path to write rejected data to; files with the following suffixes are created:
    "lang"  ... rejected since the content is in an unwanted recognised language
    "mixed" ... rejected since the content is below the threshold of discerning languages
                (too similar to two top scoring recognised languages)
    "small" ... rejected since the content is too short to reliably determine its language

LANG_RATIO_THR threshold ratio of wordlist scores of top/second most scoring languages
    (NONE not to filter mixed language at all)
    1.1 is ok for different languages, 1.01 is better for very close languages

Input format: Vertical (tokenised text, one token per line), <doc/> and <p/> structures.

Output format: The same as input with the following additions:
    Attribute doc.lang -- the language with the top score
    Attribute doc.lang_scores -- scores for each recognised language
    Structure par_langs with attributes lang and lang_scores -- the same for paragraphs
    New token attributes (columns): scores for recognised languages (command line order)

A word form is expected in the first column of the input.
The script keeps any other token attributes (e.g. lemma, tag).
The wordlist comparison is case insensitive.
'''

import io, re, sys
from itertools import chain
from math import log

if len(sys.argv) < 8 or '-h' in sys.argv or '--help' in sys.argv:
    sys.stderr.write(USAGE)
    sys.exit(1)

langs = list(map(str.lower, sys.argv[1:-3:2]))
wl_paths = sys.argv[2:-3:2]
langs_to_accept = sys.argv[-3].lower().split(',')
rejected_docs_file = sys.argv[-2]
if sys.argv[-1] == 'NONE':
    LANG_RATIO_THRESHOLD = 0.0
else:
    LANG_RATIO_THRESHOLD = float(sys.argv[-1])
if 'all' not in langs_to_accept:
    for lang in langs_to_accept:
        if lang not in langs:
            raise ValueError('%s is not in the list of languages to recognise.' % lang)
if '.gz' in ''.join(wl_paths):
    from gzip import open as gzip_open
if '.xz' in ''.join(wl_paths):
    from lzma import open as lzma_open

#Additional filtering is applied since comparing structures too short is not reliable and we
#do not want that in the corpus anyway. Ideally, there should be a separate script for that.
MIN_WORD_LEN = 1
MAX_WORD_LEN = 20
MIN_DOC_LINES = 20 #throw away documents shorter than this count of lines
MIN_DOC_GOOD_WORDS = 10 #throw away documents shorter than this count of tokens in good paragraphs
MIN_DOC_GOOD_WORD_RATIO = 0.25 #min ratio of non-rare dict words in the target language in the doc
MIN_PAR_LINES = 5 #account paragraphs of at least this count of lines in the document size
MIN_PAR_GOOD_WORDS = 1 #account paragraphs of at least this count of wordlist tokens in the document size
MIN_GOOD_WORD_RELFRQ = 1.0 #log(token_frq * norm_wl_size, 10) >= MIN_GOOD_WORD_RELFRQ => good word

stdin = io.open(sys.stdin.fileno(), mode='rt', encoding='utf-8', errors='ignore',
    buffering=10000000, newline='\n')
stdout = io.open(sys.stdout.fileno(), mode='wt', encoding='utf-8', buffering=10000000)

tag_re = re.compile('<(/?)([^ >]+)[ >]')
par_tag_re = re.compile('</?p[ >]')

#read language wordlists
wls = {}
word_re = re.compile('.{%d,%d}$' % (MIN_WORD_LEN, MAX_WORD_LEN))
for lang, wl_path in zip(langs, wl_paths):
    sys.stderr.write('Reading %s frequency wordlist\n' % lang)
    #read words with frequencies
    lang_wl_tmp, wl_size = {}, 0
    if wl_path.endswith('.xz'):
        file_wrapper = lzma_open
    elif wl_path.endswith('.gz'):
        file_wrapper = gzip_open
    else:
        file_wrapper = io.open
    with file_wrapper(wl_path, mode='rt', encoding='utf-8', newline='\n') as fp:
        for line in fp.readlines():
            word, freq = line.split('\t')
            word = word.lower()
            if word_re.match(word):
                freq = int(freq)
                wl_size += freq
                lang_wl_tmp[word] = lang_wl_tmp.get(word, 0) + freq
    #recalculate to relative frequencies
    norm_wl_size = 1000000000.0 / wl_size
    wls[lang] = dict([(word, log(freq * norm_wl_size, 10)) for word, freq in lang_wl_tmp.items()])
    sys.stderr.write('%s wordlist read: %d items\n' % (lang, len(wls[lang])))
sys.stderr.write('All wordlists read\n')

def read_big_structures(fp, structure_tag, buffer_size=20000000):
    structure_start_re = re.compile('^<%s ' % structure_tag, re.M)
    buffer = ''
    while True:
        new_data = fp.read(buffer_size)
        if not new_data:
            break
        buffer += new_data
        starting_positions = [m.start() for m in structure_start_re.finditer(buffer)]
        if starting_positions == []:
            continue
        for i in range(len(starting_positions) - 1):
            start = starting_positions[i]
            end = starting_positions[i + 1]
            yield buffer[start:end]
        buffer = buffer[starting_positions[-1]:]
    if buffer != '':
        yield buffer

def get_highest_lang(lang_scores):
    #compare score ratios, select the language with the highest score if ratios > threshold
    result_lang, result_lang_score = None, 0.0
    for lang, score in lang_scores.items():
        if score > result_lang_score:
            result_lang, result_lang_score = lang, score
    over_threshold = True
    for lang, score in lang_scores.items():
        if LANG_RATIO_THRESHOLD and lang != result_lang and score > 0.0 \
                and result_lang_score / score < LANG_RATIO_THRESHOLD:
            over_threshold = False #cannot determine the language
            break
    if not result_lang or not over_threshold:
        result_lang = 'mixed'
        #result_lang = 'czech' ## ugly fix 
    return result_lang

#filter documents
docs_read, doc_count, par_count = 0, {}, {}
with io.open(rejected_docs_file + '_lang', mode='wt', encoding='utf-8') as fp_rejected, \
        io.open(rejected_docs_file + '_mixed', mode='wt', encoding='utf-8') as fp_mixed, \
        io.open(rejected_docs_file + '_small', mode='wt', encoding='utf-8') as fp_small_size:
    doc_id = 0
    for doc in read_big_structures(stdin, 'doc'):
        doc_header, doc_rest = doc.split('\n', 1)
        if not '">' in doc_header:
            sys.stderr.write('\nWARNING: skipping invalid doc header: "%s"\n' % doc_header)
            continue
        docs_read += 1
        #sum wordlist relative frequencies per a language for all words in the document
        doc_lines, doc_size_by_lang = [], dict((lang, 0) for lang in langs)
        for line in filter(None, doc_rest.split('\n')):
            if line[0] != '<':
                word = line.split('\t')[0].strip().lower()
                word_scores = {}
                for lang in langs:
                    lang_score = wls[lang].get(word, 0.0)
                    word_scores[lang] = lang_score
                doc_lines.append((word, word_scores, line.rstrip()))
            elif par_tag_re.match(line) and doc_lines and doc_lines[-1][1] != True:
                #start or end of paragraph and not on the previous line
                doc_lines.append((None, True, line.rstrip())) #mark end of paragraph
            else:
                doc_lines.append((None, False, line.rstrip())) #other tag
        #split paragraphs of the document by language
        doc_lines_all = []
        doc_lines_by_lang = dict((lang, []) for lang in langs)
        par_lines, par_words = [], []
        par_scores, par_size = dict((lang, 0.0) for lang in langs), 0
        doc_scores_langs = dict((lang, dict(
            (lang, 0.0) for lang in langs)) for lang in langs) #lang scores for doc parts by lang
        for word, word_info, line in doc_lines:
            if word and word[0] != '<':
                par_lines.append('%s\t%s' % (line, '\t'.join(
                    '%.2f' % word_info[lang] for lang in langs)))
                par_words.append(word)
                for lang in langs:
                    par_scores[lang] += word_info[lang]
                par_size += 1
            else:
                par_lines.append(line)
                if word_info: #end of paragraph
                    par_lang = get_highest_lang(par_scores)
                    score_info = ', '.join(['%s: %.2f' %
                        (lang, par_scores[lang]) for lang in langs])
                    if par_lang == 'mixed':
                        '''
                        fp_mixed.write('<par_langs lang="%s" lang_scores="%s"/>\n%s\n' %
                            (par_lang, score_info, '\n'.join(par_lines)))
                        '''
                        # doc_lines_all.append('<par_langs lang="%s" lang_scores="%s"/>' % (par_lang, score_info)) ## added
                        doc_lines_all.extend(par_lines) ## added

                    else:
                        par_good_word_count = len([w for w in par_words
                            if wls[par_lang].get(w, 0.0) >= MIN_GOOD_WORD_RELFRQ])
                        if par_size >= MIN_PAR_LINES and par_good_word_count >= MIN_PAR_GOOD_WORDS:
                            doc_size_by_lang[par_lang] += par_good_word_count
                        '''
                        doc_lines_by_lang[par_lang].append(
                            '<par_langs lang="%s" lang_scores="%s"/>' % (par_lang, score_info))
                        doc_lines_by_lang[par_lang].extend(par_lines)
                        '''
                        # doc_lines_all.append('<par_langs lang="%s" lang_scores="%s"/>' % (par_lang, score_info)) ## added
                        doc_lines_all.extend(par_lines) ## added
                        par_count[par_lang] = par_count.get(par_lang, 0) + 1
                        for lang in langs:
                            doc_scores_langs[par_lang][lang] += par_scores[lang]
                    par_lines, par_words = [], []
                    par_scores, par_size = dict((lang, 0.0) for lang in langs), 0
        ## doc_header1 = doc_header.replace(' lang="', ' lang_old="') ## added
        print(doc_header) ## added
        for x in doc_lines_all: ## added
            print(x)            ## added
        print('</doc>')
        #write out docs by target lang
        '''
        for target_lang in langs:
            if doc_lines_by_lang[target_lang]:
                doc_header1 = doc_header.replace(' lang="', ' lang_old="') #rename the lang attr
                score_info = ', '.join(['%s: %.2f' %
                    (lang, doc_scores_langs[target_lang][lang]) for lang in langs])
                doc_header1 = doc_header1.replace('">', '" lang="%s" lang_scores="%s">' %
                    (target_lang, score_info))
                doc_lines_by_lang_count = len(doc_lines_by_lang[target_lang])
                if doc_lines_by_lang_count < MIN_DOC_LINES \
                        or doc_size_by_lang[target_lang] < MIN_DOC_GOOD_WORDS \
                        or doc_size_by_lang[target_lang] / doc_lines_by_lang_count \
                        < MIN_DOC_GOOD_WORD_RATIO:
                    fp_out = fp_small_size # ugly fix
                    #fp_out = stdout
                    ###
                    doc_count['small'] = doc_count.get('small', 0) + 1
                else:
                    if 'all' in langs_to_accept or target_lang in langs_to_accept:
                        fp_out = stdout
                    else:
                        fp_out = fp_rejected
                    doc_count[target_lang] = doc_count.get(target_lang, 0) + 1
                fp_out.write('%s\n%s\n</doc>\n' %
                    (doc_header1, '\n'.join(doc_lines_by_lang[target_lang])))
        doc_id += 1
        if doc_id % 1000 == 0:
            docs_kept = sum(doc_count.get(x, 0) for x in langs)
            sys.stderr.write('\r%d/%d docs written' % (docs_kept, doc_id))
        '''

docs_kept = sum(doc_count.get(x, 0) for x in langs)
sys.stderr.write('\nDone, %d docs read, %d docs written\n' % (docs_read, docs_kept))
sys.stderr.write('Document counts:\n')
for lang, count in doc_count.items():
    sys.stderr.write('\t%s\t%d\n' % (lang, count))
sys.stderr.write('Paragraph counts:\n')
for lang, count in par_count.items():
    sys.stderr.write('\t%s\t%d\n' % (lang, count))

