#! /usr/bin/env python3
# -*- mode: Python; -*-

import os, re, sys, traceback, unicodedata
import hfst

from vrtargslib import trans_args, trans_main
from vrtargslib import BadData, BadCode

from vrtnamelib import xname
from vrtnamelib import namelist, nameindex, insertnames, isnames

from vrtdatalib import asrecord
from vrtdatalib import escape, unescape

# some overrides to the omorfi transducer
from libtdpalphalookup import additions, TAGCAT

from outsidelib import ALOMORFI

def parsearguments():
    description = '''

    Look up the words in a vrt document in the "particular version of
    Omorfi" that came with the alpha version of the Turku dependency
    parser. Build suitable feature sets (possible word class tags) to
    use with the MarMoT model of the parser, not entirely unlike the
    parser itself does. Insert the feature sets after the words. - The
    vrt document must have position names before any token.

'''

    parser = trans_args(description = description)

    parser.add_argument('--feats', '-f', metavar = 'name',
                        type = xname, default = 'marmot.in',
                        help = 'output field name (default marmot.in)')
    parser.add_argument('--word', '-w', metavar = 'name',
                        type = xname, default = 'word',
                        help = 'input field name (default word)')
    parser.add_argument('--debug', dest = 'debug',
                        choices = ['raw', 'bare', 'basic',
                                   'edited', 'extended'],
                        help = 'show various stages (not VRT)')

    args = parser.parse_args()
    args.prog = parser.prog
    return args

## TODO PROPERLY
# extracted from omorfi_pos.py to work out analyze_reading, assuming
# current hfst-optimized-lookup be suitable looker-up (instead of the
# Java they used) Nupdate! current Python looker-up is perfect!

# import omorfi_pos as omor

### TODO comments are from omorfi_pos.py by the Turku NLP group
'''
lemma,tags=analyze_reading(reading)
returns the lemma and the list of tags of the last compound member's
last full derivation (i.e. derivation which produces POS)
'''

#Returns the lemma and list of readings which you can then feed into analyze_tagset
#Raises ValueError if something goes haywire [[DOES NOT]]

re_tag = re.compile('<([^<>]+)>')

def raw_analyses(form):
    '''(a debug level)'''
    return [ ('_', [analysis])
             for analysis, weight in FTB.lookup(unescape(form)) ]

def bare_analyses(form):
    '''(a debug level)'''
    return [ ('_', [re.sub('@[^@]+@', '', analysis)])
             for analysis, weight in FTB.lookup(unescape(form)) ]

def basic_analyses(form):

    # punctuation and number checks are both in omorfi_lookup
    # (relevant before marmot) and then separately after calling
    # marmot in marmot_wrapper

    # all punctuation characters (following omorfi_pos.omorfi_lookup)
    # (categories are punctuation, mathematical/modifier symbol)
    if all(cat.startswith('P') or cat in ('Sm', 'Sk')
           for cat in map(unicodedata.category, form)):
        return [ (form, ['Punct']) ]

    # number-pattern match (following omorfi_pos.omorfi_lookup)
    # (codes are various dashes)
    if re.fullmatch('[0-9.,:\u2012\u2013\u2014\u2015\u2053~-]+', form):
        return [ (form, ['Num']) ]

    return [ basic_analysis(analysis)
             for analysis, weight in FTB.lookup(unescape(form)) ]

def basic_analysis(analysis):
    '''From omor.analyze_reading. Extract the more or less useful content
    of a reading.

    Work on an analysis line, and there is no such thing as an
    analysis for unrecognized input.

    reading => lemma, tags

    '''

    # could be more precise but meh
    analysis = re.sub('@[^@]+@', '', analysis)

    if analysis.startswith('+<'):
        # +	+<Punct>
        parts = [analysis]
    else:
        # NOTE form not included in analysis!
        # Heinäpaalin	heinä<N><Sg><Nom><Cmpnd>+paali<N><Sg><Gen><Cap>
        # Heinäpaalin	heinä+#paali<N><Sg><Gen><Cap>
        # !! looks like + is never a character in another recognized token
        # !! and so does not occur in analyses - not sure, though
        parts = analysis.split('+')

    lemmaparts = []
    for part in parts:
        if part.startswith('#') and part not in ('#', '#<Punct>'):
            # #paali<N><Sg><Gen><Cap>
            part = part[1:]
        m = re_tag.search(part)
        lemmaparts.append(part if m is None else part[:m.start()])
    lemma = '|'.join(lemmaparts)

    # from the last part: [N, Sg, Gen, Cap]
    tags = re_tag.findall(parts[-1])

    # "actual" tags down to derivation
    # NOTE form not included in analysis
    # kierittäminen	kieriä<V><Der_ttaa><V><Der_minen><N><Sg><Nom><cap>
    # kierittäminen	kierittää<V><Der_minen><N><Sg><Nom><cap>

    pos = 0
    for k, tag in enumerate(reversed(tags), start = 1 - len(tags)):
        if TAGCAT[tag] == 'POS':
            if pos == 0:
                pos = -k
            else:
                # unknown whether this ever happens - was an assertion
                # ALERT broken now that not working in bytes any more
                sys.stderr.buffer(b'testing:warning:another pos\n')
                sys.stderr.buffer(b'testing:warning:' + reading + b'\n')
                return lemma, tags[pos:]
        elif pos > 0 and TAGCAT[tag] == 'DRV':
            return lemma, tags[-k:]
        else:
            # they would also log some issues but let them be:
            # - anything other than DRV before POS
            # - other DRV than Der_u, Der_minen after POS
            pass
    else:
        return lemma, tags

def edited_analyses(form):
    return tuple((lemma, edited_tags(form, tags))
                 for lemma, tags in basic_analyses(form))

def edited_tags(form, tags):
    '''From omorfi_postprocess. Then remains to add additional readings to
    the cohort for CC/CS and for some particular tokens.

    '''
    tags = [ ('Adv' if tag == 'Pcle' else
              'Px3' if tag in ('PxSg3', 'PxPl3') else tag)
             for tag in tags ]

    # fix this here so it will not cause any more problems when there
    # need be a POS later (CC, CS are SUBCAT so there is no POS, and
    # Para is whatever - some two-dot thing? seen it now anyway) -
    # they would map these only at the point where they need a POS
    tags and tags[0] in ('CC', 'CS') and tags.insert(0, 'C')
    tags and tags[0] == 'Para' and tags.insert(0, 'Punct')

    if tags and tags[-1] in ('Cap', 'cap', 'CAP'): tags.pop()
    if form[:1].isupper(): tags.append('Up')
    return tuple(tags)

def extended_analyses(form):
    return extended_cohort(form, set(edited_analyses(form)))

def extended_cohort(form, analyses):
    extCC = { (lemma, tuple(('CC' if tag == 'CS' else tag) for tag in tags))
              for lemma, tags in analyses if 'CS' in tags }
    extCS = { (lemma, tuple(('CS' if tag == 'CC' else tag) for tag in tags))
              for lemma, tags in analyses if 'CC' in tags }
    ext = additions.get(form, set())

    # not happy with no POS so throw such out if good ones remain
    # (there are bad readings in "sun" / "sunkin" / "sunkaan") wait no
    # that is not right must introduce C as POS for CC, CS first, as
    # is done now in the edited stage
    bad = {
        (lemma, tags)
        for lemma, tags in analyses
        if not any(TAGCAT[tag] == 'POS' for tag in tags)
    }

    return (analyses | extCC | extCS | ext) - bad or analyses

def marmot_feats(form):
    '''All else is overkill, actually.'''

    analyses = extended_analyses(form)

    poses = (next((tag for tag in tags if TAGCAT[tag] == 'POS'), None)
             for lemma, tags in analyses)
    feats = sorted(set(filter(None, poses)))

    return'#'.join('POS_' + feat for feat in feats) or '_'

def main(args, ins, ous):
    # no reason to stay in binary here that ftb.lookup uses unicode
    # but should ensure that stdin/stdout are UTF-8 and how does one
    # do that again?

    global FTB
    # that be optimized-lookup form transducer (specific transducer!)
    tin = hfst.HfstInputStream(ALOMORFI)
    FTB = next(tin)
    tin.close() # it *should* be made a context manager TOREPORT

    # act on names and tokens, pass other lines on as they are
    pos = None
    for line in ins:
        if isnames(line):
            pos = nameindex(namelist(line), args.word)
            print(insertnames(line, args.word, args.feats),
                  end = '', file = ous)
        elif line.isspace() or line.startswith('<'):
            print(line, end = '', file = ous)
        elif pos is None:
            raise BadData('no names before data')
        elif args.debug is None:
            record = asrecord(line)
            record.insert(pos + 1, escape(marmot_feats(unescape(record[pos]))))
            print(*record, sep = '\t', file = ous)
        else:
            record = asrecord(line)
            form = record[pos]
            analyses = (raw_analyses if args.debug == 'raw' else
                        bare_analyses if args.debug == 'bare' else
                        basic_analyses if args.debug == 'basic' else
                        edited_analyses if args.debug == 'edited' else
                        extended_analyses)(form)
            if not analyses:
                print(form, form, '_', sep = '\t', file = ous)
                print(file = ous)
            else:
                for lemma, tags in analyses:
                    print(form, lemma, '|'.join(tags), sep = '\t', file = ous)
                else:
                    print(file = ous)

if __name__ == '__main__':
    trans_main(parsearguments(), main)
