#!/usr/bin/env python3

import argparse
import json
import re
import typing

from seqeval.metrics import performance_measure

from bunkai.constant import METACHAR_LINE_BREAK, METACHAR_SENTENCE_BOUNDARY

REGEXP_SB_WITH_BLANKS_SPAN: str = f'({METACHAR_SENTENCE_BOUNDARY})([\\s{METACHAR_LINE_BREAK}]*)'
RE_SB_WITH_BLANKS_SPAN = re.compile(REGEXP_SB_WITH_BLANKS_SPAN)
RE_SBS_SPAN = re.compile(f'{METACHAR_SENTENCE_BOUNDARY}+')


def get_bo(text: str, lbonly: bool) -> typing.List[str]:
    """Get a list of BIO sequence from input text."""
    assert METACHAR_SENTENCE_BOUNDARY + METACHAR_SENTENCE_BOUNDARY not in text
    assert not text.startswith(METACHAR_SENTENCE_BOUNDARY)
    ret: typing.List[str] = []
    exist_under_bar: bool = False
    for char in text:
        if char == METACHAR_SENTENCE_BOUNDARY:
            if lbonly and exist_under_bar is False:
                continue
            else:
                ret[-1] = 'SEP'
                exist_under_bar = False
                continue
        if lbonly:
            if char == METACHAR_LINE_BREAK:
                exist_under_bar = True
                ret.append('O')
        else:
            ret.append('O')
    return ret


def get_score(data: typing.Dict[str, int]) -> typing.Dict[str, float]:
    ret: typing.Dict[str, float] = {}
    try:
        precision = data['TP'] / float(data['TP'] + data['FP'])
    except ZeroDivisionError:
        precision = float('nan')
    recall = data['TP'] / float(data['TP'] + data['FN'])
    ret['f1'] = 2 * precision * recall / (recall + precision)
    ret['precision'] = precision
    ret['recall'] = recall
    return ret


def evaluate(golds: typing.List[str], systems: typing.List[str], lbonly: bool) -> str:
    assert len(golds) == len(systems)
    golds_bo = []
    systems_bo = []
    for gold, system in zip(golds, systems):
        assert gold.replace(METACHAR_SENTENCE_BOUNDARY, '') == system.replace(METACHAR_SENTENCE_BOUNDARY, ''), \
            f'{gold}\n{system}'

        system_bo = get_bo(system, lbonly)
        systems_bo.append(system_bo)

        gold_bo = get_bo(gold, lbonly)
        golds_bo.append(gold_bo)
        assert len(gold_bo) == len(system_bo), f'gold={gold} \nsystem={system}\n' \
                                               f'N(gold)={len(gold_bo)} N(systems)={len(system_bo)}'

    measure: typing.Dict[str, int] = performance_measure(golds_bo, systems_bo)
    score = get_score(measure)
    return f'{json.dumps(measure, sort_keys=True)}\n{json.dumps(score, sort_keys=True)}\n'


def trim(val: str) -> str:
    old: str = val
    new: str = val
    while True:
        new = RE_SBS_SPAN.sub(METACHAR_SENTENCE_BOUNDARY, RE_SB_WITH_BLANKS_SPAN.sub(r'\2\1', old))
        if old == new:
            break
        old = new
    return new


def get_opts() -> argparse.Namespace:
    oparser = argparse.ArgumentParser()
    oparser.add_argument("--input", "-i", type=argparse.FileType('r'), required=True)
    oparser.add_argument("--gold", "-g", type=argparse.FileType('r'), required=False)
    oparser.add_argument("--output", "-o", type=argparse.FileType('w'), default='-')
    oparser.add_argument("--lb", action="store_true", help='Checks only Line-break marked with "\u2502"')
    oparser.add_argument("--trim", action="store_true")
    return oparser.parse_args()


def main() -> None:
    opts = get_opts()
    with opts.output as outf:
        if opts.trim:
            for line in opts.input:
                outf.write(trim(line[:-1]))
                outf.write('\n')
        else:
            r = evaluate([l1[:-1] for l1 in opts.gold.readlines()],
                         [m[:-1] for m in opts.input.readlines()], opts.lb)
            outf.write(r)


if __name__ == '__main__':
    main()
