from builtins import range, str

from fonduer.candidates.models.span import TemporarySpan
from fonduer.utils.config import get_config
from fonduer.utils.data_model_utils import (
    get_cell_ngrams,
    get_col_ngrams,
    get_head_ngrams,
    get_row_ngrams,
)
from fonduer.utils.utils_table import min_col_diff, min_row_diff, num_cols, num_rows

FEAT_PRE = "TAB_"
DEF_VALUE = 1

unary_tablelib_feats = {}
binary_tablelib_feats = {}

settings = get_config()


def get_table_feats(candidates):
    candidates = candidates if isinstance(candidates, list) else [candidates]
    for candidate in candidates:
        args = tuple([m.span for m in candidate.get_mentions()])
        if not (isinstance(args[0], TemporarySpan)):
            raise ValueError(
                "Accepts Span-type arguments, %s-type found." % type(candidate)
            )

        # Unary candidates
        if len(args) == 1:
            span = args[0]
            if span.stable_id not in unary_tablelib_feats:
                unary_tablelib_feats[span.stable_id] = set()
                for f, v in tablelib_unary_features(span):
                    unary_tablelib_feats[span.stable_id].add((f, v))

            for f, v in unary_tablelib_feats[span.stable_id]:
                yield candidate.id, FEAT_PRE + f, v

        # Binary candidates
        elif len(args) == 2:
            span1, span2 = args
            if span1.sentence.is_tabular() or span2.sentence.is_tabular():
                for span, pre in [(span1, "e1_"), (span2, "e2_")]:
                    if span.stable_id not in unary_tablelib_feats:
                        unary_tablelib_feats[span.stable_id] = set()
                        for f, v in tablelib_unary_features(span):
                            unary_tablelib_feats[span.stable_id].add((f, v))

                    for f, v in unary_tablelib_feats[span.stable_id]:
                        yield candidate.id, FEAT_PRE + pre + f, v

                if candidate.id not in binary_tablelib_feats:
                    binary_tablelib_feats[candidate.id] = set()
                    for f, v in tablelib_binary_features(span1, span2):
                        binary_tablelib_feats[candidate.id].add((f, v))

                for f, v in binary_tablelib_feats[candidate.id]:
                    yield candidate.id, FEAT_PRE + f, v
        else:
            raise NotImplementedError(
                "Only handles unary and binary candidates currently"
            )


def tablelib_unary_features(span):
    """
    Table-/structure-related features for a single span
    """
    if not span.sentence.is_tabular():
        return
    sentence = span.sentence
    for attrib in settings["featurization"]["table"]["unary_features"]["attrib"]:
        for ngram in get_cell_ngrams(
            span,
            n_max=settings["featurization"]["table"]["unary_features"][
                "get_cell_ngrams"
            ]["max"],
            attrib=attrib,
        ):
            yield "CELL_%s_[%s]" % (attrib.upper(), ngram), DEF_VALUE
        for row_num in range(sentence.row_start, sentence.row_end + 1):
            yield "ROW_NUM_[%s]" % row_num, DEF_VALUE
        for col_num in range(sentence.col_start, sentence.col_end + 1):
            yield "COL_NUM_[%s]" % col_num, DEF_VALUE
        # NOTE: These two features could be accounted for by HTML_ATTR in
        # structural features
        yield "ROW_SPAN_[%d]" % num_rows(sentence), DEF_VALUE
        yield "COL_SPAN_[%d]" % num_cols(sentence), DEF_VALUE
        for axis in ["row", "col"]:
            for ngram in get_head_ngrams(
                span,
                axis,
                n_max=settings["featurization"]["table"]["unary_features"][
                    "get_head_ngrams"
                ]["max"],
                attrib=attrib,
            ):
                yield "%s_HEAD_%s_[%s]" % (axis.upper(), attrib.upper(), ngram), 1
        for ngram in get_row_ngrams(
            span,
            n_max=settings["featurization"]["table"]["unary_features"][
                "get_row_ngrams"
            ]["max"],
            attrib=attrib,
        ):
            yield "ROW_%s_[%s]" % (attrib.upper(), ngram), DEF_VALUE
        for ngram in get_col_ngrams(
            span,
            n_max=settings["featurization"]["table"]["unary_features"][
                "get_col_ngrams"
            ]["max"],
            attrib=attrib,
        ):
            yield "COL_%s_[%s]" % (attrib.upper(), ngram), DEF_VALUE
        # TODO:
        #  for ngram in get_row_ngrams(
        #      span, n_max=2, attrib=attrib, direct=False, infer=True
        #  ):
        #      yield "ROW_INFERRED_%s_[%s]" % (attrib.upper(), ngram), DEF_VALUE
        #  for ngram in get_col_ngrams(
        #      span, n_max=2, attrib=attrib, direct=False, infer=True
        #  ):
        #      yield "COL_INFERRED_%s_[%s]" % (attrib.upper(), ngram), DEF_VALUE


def tablelib_binary_features(span1, span2):
    """
    Table-/structure-related features for a pair of spans
    """
    binary_features = settings["featurization"]["table"]["binary_features"]
    if span1.sentence.is_tabular() and span2.sentence.is_tabular():
        if span1.sentence.table == span2.sentence.table:
            yield u"SAME_TABLE", DEF_VALUE
            if span1.sentence.cell is not None and span2.sentence.cell is not None:
                row_diff = min_row_diff(
                    span1.sentence,
                    span2.sentence,
                    absolute=binary_features["min_row_diff"]["absolute"],
                )
                col_diff = min_col_diff(
                    span1.sentence,
                    span2.sentence,
                    absolute=binary_features["min_col_diff"]["absolute"],
                )
                yield u"SAME_TABLE_ROW_DIFF_[%s]" % row_diff, DEF_VALUE
                yield u"SAME_TABLE_COL_DIFF_[%s]" % col_diff, DEF_VALUE
                yield u"SAME_TABLE_MANHATTAN_DIST_[%s]" % str(
                    abs(row_diff) + abs(col_diff)
                ), DEF_VALUE
                if span1.sentence.cell == span2.sentence.cell:
                    yield u"SAME_CELL", DEF_VALUE
                    yield u"WORD_DIFF_[%s]" % (
                        span1.get_word_start_index() - span2.get_word_start_index()
                    ), DEF_VALUE
                    yield u"CHAR_DIFF_[%s]" % (
                        span1.char_start - span2.char_start
                    ), DEF_VALUE
                    if span1.sentence == span2.sentence:
                        yield u"SAME_SENTENCE", DEF_VALUE
        else:
            if span1.sentence.cell is not None and span2.sentence.cell is not None:
                yield u"DIFF_TABLE", DEF_VALUE
                row_diff = min_row_diff(
                    span1.sentence,
                    span2.sentence,
                    absolute=binary_features["min_row_diff"]["absolute"],
                )
                col_diff = min_col_diff(
                    span1.sentence,
                    span2.sentence,
                    absolute=binary_features["min_col_diff"]["absolute"],
                )
                yield u"DIFF_TABLE_ROW_DIFF_[%s]" % row_diff, DEF_VALUE
                yield u"DIFF_TABLE_COL_DIFF_[%s]" % col_diff, DEF_VALUE
                yield u"DIFF_TABLE_MANHATTAN_DIST_[%s]" % str(
                    abs(row_diff) + abs(col_diff)
                ), DEF_VALUE
