import copy
import re
from collections import defaultdict
from dataclasses import dataclass, field
from enum import Enum, Flag, auto
from functools import cmp_to_key
from itertools import chain
from typing import Dict, Optional, Union, List, Set, Tuple, NamedTuple, Any

import networkx as nx
from loguru import logger
from tqdm import tqdm
from typing_extensions import TypedDict

from grams.algorithm.kg_index import KGObjectIndex
from grams.algorithm.literal_match import WikidataValueType, TextParser, LiteralMatcher
from grams.inputs.linked_table import LinkedTable
from kgdata.wikidata.models import QNode, DataValue, WDProperty, WDClass
import sm.misc as M
from sm.misc.graph import viz_graph


@dataclass(frozen=True)
class Span:
    start: int
    end: int

    @property
    def length(self):
        return self.end - self.start


@dataclass
class CellNode:
    id: str
    value: str
    column: int
    row: int
    qnode_ids: List[str]
    qnodes_span: Dict[str, List[Span]]

    @property
    def is_cell(self):
        return True

    @property
    def is_literal_value(self):
        return False

    @property
    def is_entity_value(self):
        return False

    @property
    def is_statement(self):
        return False


@dataclass
class ContextSpan:
    # the text of the context
    text: str
    span: Span

    def get_text_span(self):
        return self.text[self.span.start : self.span.end]


@dataclass
class LiteralValueNode:
    id: str
    value: DataValue
    # not none if it is appear in the context
    context_span: Optional[ContextSpan]

    @property
    def is_cell(self):
        return False

    @property
    def is_literal_value(self):
        return True

    @property
    def is_entity_value(self):
        return False

    @property
    def is_statement(self):
        return False

    @property
    def is_context(self):
        return self.context_span is not None


@dataclass
class EntityValueNode:
    id: str
    qnode_id: str
    # not none if it is appear in the context
    context_span: Optional[ContextSpan]

    @property
    def is_cell(self):
        return False

    @property
    def is_literal_value(self):
        return False

    @property
    def is_entity_value(self):
        return True

    @property
    def is_statement(self):
        return False

    @property
    def is_context(self):
        return self.context_span is not None


EdgeFlowSource = NamedTuple("EdgeFlowSource", [("source_id", str), ("edge_id", str)])
EdgeFlowTarget = NamedTuple("EdgeFlowTarget", [("target_id", str), ("edge_id", str)])


class LinkGenMethod(Enum):
    """Methods to generate a link"""

    # this statement is generated by exact matching from the link
    FromWikidataLink = "from_wikidata_link"
    FromLiteralMatchingFunc = "from_literal_matching_function"


@dataclass
class FlowProvenance:
    """Contain information regarding how this relationship/flow has been generated (typically coming from the matching algorithm)"""

    # method that
    gen_method: LinkGenMethod
    gen_method_arg: Any
    prob: float

    def merge(self, another: "FlowProvenance") -> Union["FlowProvenance", None]:
        """Try to merge the provenance, if we cannot merge them, just return None"""
        if self.gen_method != another.gen_method:
            return None
        if self.gen_method == LinkGenMethod.FromWikidataLink:
            return self
        if self.gen_method == LinkGenMethod.FromLiteralMatchingFunc:
            if self.gen_method_arg == another.gen_method_arg:
                if self.prob > another.prob:
                    return self
                return another
            return None
        raise M.UnreachableError()

    @staticmethod
    def merge_lst(
        lst1: List["FlowProvenance"], lst2: List["FlowProvenance"]
    ) -> List["FlowProvenance"]:
        """Assume that items within each list (lst1 & lst2) are not mergeable"""
        if len(lst1) == 0:
            return lst2
        elif len(lst2) == 0:
            return lst1

        lst = copy.copy(lst1)
        for item in lst2:
            for i in range(len(lst1)):
                resp = lst[i].merge(item)
                if resp is None:
                    # cannot merge them
                    lst.append(item)
                    break
                else:
                    # merge lst[i] & item
                    lst[i] = resp
                    break
        return lst


@dataclass
class StatementNode:
    id: str
    # id of the qnode that contains the statement
    qnode_id: str
    # predicate of the statement
    predicate: str
    # whether this statement actually exist in KG
    is_in_kg: bool

    # recording which link in the source is connected to the target.
    forward_flow: Dict[EdgeFlowSource, Set[EdgeFlowTarget]] = field(
        default_factory=dict
    )
    reversed_flow: Dict[EdgeFlowTarget, Set[EdgeFlowSource]] = field(
        default_factory=dict
    )
    flow: Dict[Tuple[EdgeFlowSource, EdgeFlowTarget], List[FlowProvenance]] = field(
        default_factory=dict
    )

    @property
    def value(self):
        return self.id

    @property
    def is_cell(self):
        return False

    @property
    def is_literal_value(self):
        return False

    @property
    def is_entity_value(self):
        return False

    @property
    def is_statement(self):
        return True

    def track_provenance(
        self,
        source_flow: EdgeFlowSource,
        target_flow: EdgeFlowTarget,
        provenances: List[FlowProvenance],
    ):
        ptr = self.forward_flow
        if source_flow not in ptr:
            ptr[source_flow] = set()
        ptr[source_flow].add(target_flow)

        ptr = self.reversed_flow
        if target_flow not in ptr:
            ptr[target_flow] = set()
        ptr[target_flow].add(source_flow)

        # TODO: merge the provenance if we can
        if (source_flow, target_flow) not in self.flow:
            self.flow[source_flow, target_flow] = provenances
        else:
            self.flow[source_flow, target_flow] = FlowProvenance.merge_lst(
                self.flow[source_flow, target_flow], provenances
            )

    def untrack_source_flow(self, source_flow: EdgeFlowSource):
        for target_flow in self.forward_flow.pop(source_flow):
            self.flow.pop((source_flow, target_flow))
            self.reversed_flow[target_flow].remove(source_flow)
            if len(self.reversed_flow[target_flow]) == 0:
                self.reversed_flow.pop(target_flow)

    def untrack_target_flow(self, target_flow: EdgeFlowTarget):
        for source_flow in self.reversed_flow.pop(target_flow):
            self.flow.pop((source_flow, target_flow))
            self.forward_flow[source_flow].remove(target_flow)
            if len(self.forward_flow[source_flow]) == 0:
                self.forward_flow.pop(source_flow)

    def untrack_flow(self, source_flow: EdgeFlowSource, target_flow: EdgeFlowTarget):
        self.flow.pop((source_flow, target_flow))
        self.forward_flow[source_flow].remove(target_flow)
        self.reversed_flow[target_flow].remove(source_flow)

        if len(self.forward_flow[source_flow]) == 0:
            self.forward_flow.pop(source_flow)
        if len(self.reversed_flow[target_flow]) == 0:
            self.reversed_flow.pop(target_flow)

    def has_source_flow(self, source_flow):
        return source_flow in self.forward_flow

    def has_target_flow(self, target_flow):
        return target_flow in self.reversed_flow

    def has_flow(self, source_flow, target_flow):
        return (source_flow, target_flow) in self.flow

    def iter_source_flow(self, target_flow: EdgeFlowTarget):
        for source_flow in self.reversed_flow[target_flow]:
            yield source_flow, self.flow[source_flow, target_flow]

    def iter_target_flow(self, source_flow: EdgeFlowSource):
        for target_flow in self.forward_flow[source_flow]:
            yield target_flow, self.flow[source_flow, target_flow]

    def get_provenance(self, source_flow: EdgeFlowSource, target_flow: EdgeFlowTarget):
        return self.flow[source_flow, target_flow]

    def get_provenance_by_edge(self, inedge: "DGEdge", outedge: "DGEdge"):
        return self.flow[
            EdgeFlowSource(inedge.source, inedge.predicate),
            EdgeFlowTarget(outedge.target, outedge.predicate),
        ]

    def is_same_flow(self, inedge: "DGEdge", outedge: "DGEdge") -> bool:
        return (
            EdgeFlowSource(inedge.source, inedge.predicate),
            EdgeFlowTarget(outedge.target, outedge.predicate),
        ) in self.flow


DGNode = Union[CellNode, LiteralValueNode, EntityValueNode, StatementNode]


@dataclass
class DGStatementID:
    qnode_id: str
    predicate: str
    statement_index: int

    def get_id(self):
        return f"stmt:{self.qnode_id}-{self.predicate}-{self.statement_index}"

    @staticmethod
    def parse_id(id: str) -> "DGStatementID":
        m = re.match(r"stmt:([^-]+)-([^-]+)-([^-]+)", id)
        assert m is not None
        qnode_id, predicate, stmt_index = [m.group(i) for i in range(1, 4)]
        stmt_index = int(stmt_index)
        return DGStatementID(qnode_id, predicate, stmt_index)


@dataclass
class DGPathNodeStatement:
    qnode_id: str
    predicate: str
    statement_index: int
    # how this DGPath statement has been matched
    provenance: FlowProvenance

    def get_id(self):
        return DGStatementID(
            self.qnode_id, self.predicate, self.statement_index
        ).get_id()

    # ############# METHODs to construct DGPathNodeStatement from provenance ##############
    @staticmethod
    def from_FromWikidataLink(qnode_id, predicate, stmt_index):
        return DGPathNodeStatement(
            qnode_id,
            predicate,
            stmt_index,
            FlowProvenance(LinkGenMethod.FromWikidataLink, None, 1.0),
        )

    @staticmethod
    def from_FromLiteralMatchingFunc(qnode_id, predicate, stmt_index, fn_args, prob):
        return DGPathNodeStatement(
            qnode_id,
            predicate,
            stmt_index,
            FlowProvenance(LinkGenMethod.FromLiteralMatchingFunc, fn_args, prob),
        )


@dataclass
class DGPathNodeQNode:
    qnode_id: str

    def get_id(self):
        return f"ent:{self.qnode_id}"


@dataclass
class DGPathNodeLiteralValue:
    value: DataValue

    def get_id(self):
        return f"val:{self.value.to_string_repr()}"


@dataclass
class DGPathExistingNode:
    id: str

    def get_id(self):
        return self.id


DGPathNode = Union[
    DGPathNodeStatement, DGPathNodeQNode, DGPathNodeLiteralValue, DGPathExistingNode
]


@dataclass
class DGPathEdge:
    value: str
    is_qualifier: bool

    @staticmethod
    def p(value: str):
        return DGPathEdge(value, is_qualifier=False)

    @staticmethod
    def q(value: str):
        return DGPathEdge(value, is_qualifier=True)


@dataclass
class DGPath:
    # a sequence of path, is always
    sequence: List[Union[DGPathEdge, DGPathNode]] = field(default_factory=list)


@dataclass
class DGEdge:
    source: str
    target: str
    predicate: str
    is_qualifier: bool
    # deprecated, will be replaced by the information stored directly in the statement edge flow
    # paths: List[DGPath]
    is_inferred: Optional[bool] = None

    @staticmethod
    def can_link(source: DGNode, target: DGNode):
        """Test if there can be two links between the two nodes. Basically, it just enforce the constraint that no link
        cross different rows
        """
        if not source.is_cell or not target.is_cell:
            return True
        return source.row == target.row


@dataclass
class InferredNewProp:
    # information to identify the statement for the new prop (whether to reuse it, or create a new one)
    # the qnode of the new statement (the one will have this property)
    qnode_id: str
    # the new property we inferred
    new_prop: str
    # the value associated with the property (use it to compare if the statement exists) - the reason we use
    # the value instead of the whole statement is that sometimes, desire them to match the qualifier does not
    # sense, if the value exist, then some how the algorithm is already match it so it should be okay
    value: DataValue

    # the source nodes in the data graph
    source_id: str
    # the target node in the data graph that will contain the value of the prop, not the statement node
    target_id: str

    # edges that contain the qualifiers that we want to be copied to the new statement as well
    # the flow of the qualifier can be retrieve
    qualifier_edges: List[DGEdge]

    # flow provenance of the property.
    flow_provenances: List[FlowProvenance]


@dataclass
class InferredNewQualifier:
    # the statement node in the data graph that we will add the qualifier to
    statement_id: str
    # the qualifier we are going to add
    new_qualifier: str

    # source id (node that has the statement id)
    source_id: str
    # the property of the statement (the one that connect source and statement)
    property: str
    # the target node of the qualifier
    target_id: str

    # flow provenance of the qualifier
    flow_provenances: List[FlowProvenance]


class BuildDGOption(Flag):
    # disable kg search for the same entity
    DISABLE_SAME_ENT_SEARCH = auto()
    # removing unnecessary entities to keep the graph size reasonable
    PRUNING_REDUNDANT_ENT = auto()
    # disable path discovering using index
    NO_KG_INDEX = auto()


def kg_path_discovering(
    options: BuildDGOption,
    qnodes: Dict[str, QNode],
    kg_object_index: KGObjectIndex,
    dg: nx.MultiDiGraph,
    u: DGNode,
    v: DGNode,
    max_n_hop: int = 2,
    bidirection: bool = True,
):
    """Find all paths between two nodes in the graph.

    Parameters
    ----------
    options: options when building DG
    qnodes: dictionary of qnodes
    dg: the input data graph
    u: node
    v: node
    max_n_hop: hop start at 1 (current node), n_hop = 2 mean the path can go through an hidden node
    bidirection: set to false telling if we want to just find paths going from u -> v.
        This parameter just for re-using code, as we will call the function changing the order of u & v

    Returns
    -------
    """

    # noinspection PyShadowingNames
    def _path_object_search_v1(
        qnodes: Dict[str, QNode],
        kg_object_index: KGObjectIndex,
        source: QNode,
        target: QNode,
        max_n_hop=2,
    ):
        # max_n_hop = 2 mean we will find a path that go from source to target through an hidden node
        # hop start at 1 (current node)

        if options & BuildDGOption.DISABLE_SAME_ENT_SEARCH:
            if source.id == target.id:
                return []

        matches = []
        iter = source.props.items()

        for p, stmts in iter:
            if p == "P31":
                # no need to search in the instanceOf property, as the ontology is removed from the databased as they are huge
                continue

            for stmt_i, stmt in enumerate(stmts):
                # add the cell id so that we have different statement for different cell.
                has_stmt_value = False

                # to simplify the design, we do not consider a statement that its value do not exist in KQ
                # due to an error on KG
                if stmt.value.is_qnode():
                    if stmt.value.as_entity_id() not in qnodes:
                        # this can happen due to some of the qnodes is in the link, but is missing in the KG
                        # this is very rare so we can employ some check to make sure this is not due to
                        # our wikidata subset
                        is_error_in_kg = any(
                            any(
                                _s.value.is_qnode()
                                and _s.value.as_entity_id() in qnodes
                                for _s in _stmts
                            )
                            for _p, _stmts in source.props.items()
                        ) or stmt.value.as_entity_id().startswith("L")
                        if not is_error_in_kg:
                            raise Exception(
                                f"Missing qnodes in your KG subset: {stmt.value.as_entity_id()}"
                            )
                        continue

                if stmt.value.is_qnode():
                    # found by match entity
                    if stmt.value.as_entity_id() == target.id:
                        matches.append(
                            DGPath(
                                sequence=[
                                    DGPathEdge.p(p),
                                    DGPathNodeStatement.from_FromWikidataLink(
                                        source.id, p, stmt_i
                                    ),
                                    DGPathEdge.p(p),
                                ]
                            )
                        )
                        has_stmt_value = True
                    elif max_n_hop > 1:
                        stmt_value_qnode_id = stmt.value.as_entity_id()
                        if stmt_value_qnode_id.startswith("L"):
                            assert (
                                stmt_value_qnode_id not in qnodes
                            ), "The L nodes (lexical) is not in the Wikidata dump"
                            continue

                        for nextpath in _path_object_search(
                            qnodes,
                            kg_object_index,
                            qnodes[stmt_value_qnode_id],
                            target,
                            max_n_hop - 1,
                        ):
                            matches.append(
                                DGPath(
                                    sequence=[
                                        DGPathEdge.p(p),
                                        DGPathNodeStatement.from_FromWikidataLink(
                                            source.id, p, stmt_i
                                        ),
                                        DGPathEdge.p(p),
                                        DGPathNodeQNode(stmt_value_qnode_id),
                                    ]
                                    + nextpath.sequence
                                )
                            )
                            has_stmt_value = True

                for q, qvals in stmt.qualifiers.items():
                    for qval in qvals:
                        if qval.is_qnode():
                            if qval.as_entity_id() == target.id:
                                if not has_stmt_value:
                                    if stmt.value.is_qnode():
                                        pn_stmt_value = DGPathNodeQNode(
                                            stmt.value.as_entity_id()
                                        )
                                    else:
                                        pn_stmt_value = DGPathNodeLiteralValue(
                                            stmt.value
                                        )
                                    matches.append(
                                        DGPath(
                                            sequence=[
                                                DGPathEdge.p(p),
                                                DGPathNodeStatement.from_FromWikidataLink(
                                                    source.id, p, stmt_i
                                                ),
                                                DGPathEdge.p(p),
                                                pn_stmt_value,
                                            ]
                                        )
                                    )
                                    has_stmt_value = True

                                matches.append(
                                    DGPath(
                                        sequence=[
                                            DGPathEdge.p(p),
                                            DGPathNodeStatement.from_FromWikidataLink(
                                                source.id, p, stmt_i
                                            ),
                                            DGPathEdge.q(q),
                                        ]
                                    )
                                )
                            elif max_n_hop > 1:
                                qval_qnode_id = qval.as_entity_id()
                                if qval_qnode_id.startswith("L"):
                                    assert (
                                        qval_qnode_id not in qnodes
                                    ), "The L nodes (lexical) is not in the Wikidata dump"
                                    continue

                                if qval_qnode_id not in qnodes:
                                    # this can happen due to some of the qnodes is in the link, but is missing in the KG
                                    # this is very rare so we can employ some check to make sure this is not due to
                                    # our wikidata subset
                                    is_error_in_kg = any(
                                        any(
                                            _s.value.is_qnode()
                                            and _s.value.as_entity_id() in qnodes
                                            for _s in _stmts
                                        )
                                        for _p, _stmts in source.props.items()
                                    ) or qval_qnode_id.startswith("L")
                                    if not is_error_in_kg:
                                        raise Exception(
                                            f"Missing qnodes in your KG subset: {qval_qnode_id}"
                                        )
                                    continue

                                _n_matches = len(matches)
                                for nextpath in _path_object_search(
                                    qnodes,
                                    kg_object_index,
                                    qnodes[qval_qnode_id],
                                    target,
                                    max_n_hop - 1,
                                ):
                                    matches.append(
                                        DGPath(
                                            sequence=[
                                                DGPathEdge.p(p),
                                                DGPathNodeStatement.from_FromWikidataLink(
                                                    source.id, p, stmt_i
                                                ),
                                                DGPathEdge.q(q),
                                                DGPathNodeQNode(qval.as_entity_id()),
                                            ]
                                            + nextpath.sequence
                                        )
                                    )

                                if len(matches) > _n_matches and not has_stmt_value:
                                    if stmt.value.is_qnode():
                                        pn_stmt_value = DGPathNodeQNode(
                                            stmt.value.as_entity_id()
                                        )
                                    else:
                                        pn_stmt_value = DGPathNodeLiteralValue(
                                            stmt.value
                                        )
                                    matches.append(
                                        DGPath(
                                            sequence=[
                                                DGPathEdge.p(p),
                                                DGPathNodeStatement.from_FromWikidataLink(
                                                    source.id, p, stmt_i
                                                ),
                                                DGPathEdge.p(p),
                                                pn_stmt_value,
                                            ]
                                        )
                                    )
                                    has_stmt_value = True

        return matches

    def _path_object_search_v2(
        qnode: Dict[str, QNode],
        kg_object_index: KGObjectIndex,
        source: QNode,
        target: QNode,
        max_n_hop=2,
    ):
        # max_n_hop = 2 mean we will find a path that go from source to target through an hidden node
        # hop start at 1 (current node)

        if options & BuildDGOption.DISABLE_SAME_ENT_SEARCH:
            if source.id == target.id:
                return []

        matches = []
        for hop1_path in kg_object_index.iter_hop1_props(source.id, target.id):
            stmt_i = hop1_path.statement_index
            rel = hop1_path.relationship
            stmt = source.props[rel.prop][stmt_i]

            if len(rel.quals) > 0 and not rel.both:
                # the prop doesn't match, have to add it
                if stmt.value.is_qnode():
                    pn_stmt_value = DGPathNodeQNode(stmt.value.as_entity_id())
                else:
                    pn_stmt_value = DGPathNodeLiteralValue(stmt.value)
                matches.append(
                    DGPath(
                        sequence=[
                            DGPathEdge.p(rel.prop),
                            DGPathNodeStatement.from_FromWikidataLink(
                                source.id, rel.prop, stmt_i
                            ),
                            DGPathEdge.p(rel.prop),
                            pn_stmt_value,
                        ]
                    )
                )
            else:
                # the prop match
                matches.append(
                    DGPath(
                        sequence=[
                            DGPathEdge.p(rel.prop),
                            DGPathNodeStatement.from_FromWikidataLink(
                                source.id, rel.prop, stmt_i
                            ),
                            DGPathEdge.p(rel.prop),
                        ]
                    )
                )
            for qual in rel.quals:
                matches.append(
                    DGPath(
                        sequence=[
                            DGPathEdge.p(rel.prop),
                            DGPathNodeStatement.from_FromWikidataLink(
                                source.id, rel.prop, stmt_i
                            ),
                            DGPathEdge.q(qual),
                        ]
                    )
                )

        if max_n_hop > 1:
            # TODO: need to refactor the current method of searching for paths
            # as we may generate many duplicated segments
            for match_item in kg_object_index.iter_hop2_props(source.id, target.id):
                if match_item.qnode == target.id:
                    continue
                hop1_paths = match_item.hop1
                hop2_paths = match_item.hop2
                middle_qnode = qnodes[match_item.qnode]

                # we don't care about transitive here, so we do a cross product
                hop1_seqs = []
                for hop1_path in hop1_paths:
                    rel = hop1_path.relationship
                    stmt_i = hop1_path.statement_index
                    stmt = source.props[rel.prop][stmt_i]

                    if len(rel.quals) > 0 and not rel.both:
                        # we don't have the statement value yet, add it to the matches
                        # the prop doesn't match, have to add it, we don't worry about duplication
                        # as it is resolve during the merge provenance phase
                        if stmt.value.is_qnode():
                            pn_stmt_value = DGPathNodeQNode(stmt.value.as_entity_id())
                        else:
                            pn_stmt_value = DGPathNodeLiteralValue(stmt.value)
                        matches.append(
                            DGPath(
                                sequence=[
                                    DGPathEdge.p(rel.prop),
                                    DGPathNodeStatement.from_FromWikidataLink(
                                        source.id, rel.prop, stmt_i
                                    ),
                                    DGPathEdge.p(rel.prop),
                                    pn_stmt_value,
                                ]
                            )
                        )
                    else:
                        # add prop to the seqs that we need to expand next, and so stmt.value must be a qnode
                        # as it is the middle qnode
                        assert stmt.value.as_entity_id() == middle_qnode.id
                        hop1_seqs.append(
                            [
                                DGPathEdge.p(rel.prop),
                                DGPathNodeStatement.from_FromWikidataLink(
                                    source.id, rel.prop, stmt_i
                                ),
                                DGPathEdge.p(rel.prop),
                                DGPathNodeQNode(middle_qnode.id),
                            ]
                        )

                    for qual in rel.quals:
                        hop1_seqs.append(
                            [
                                DGPathEdge.p(rel.prop),
                                DGPathNodeStatement.from_FromWikidataLink(
                                    source.id, rel.prop, stmt_i
                                ),
                                DGPathEdge.p(qual),
                                DGPathNodeQNode(middle_qnode.id),
                            ]
                        )

                for hop2_path in hop2_paths:
                    rel = hop2_path.relationship
                    stmt_i = hop2_path.statement_index
                    stmt = middle_qnode.props[rel.prop][stmt_i]

                    if len(rel.quals) > 0 and not rel.both:
                        if stmt.value.is_qnode():
                            pn_stmt_value = DGPathNodeQNode(stmt.value.as_entity_id())
                        else:
                            pn_stmt_value = DGPathNodeLiteralValue(stmt.value)

                        hop2_seq = [
                            DGPathEdge.p(rel.prop),
                            DGPathNodeStatement.from_FromWikidataLink(
                                middle_qnode.id, rel.prop, stmt_i
                            ),
                            DGPathEdge.p(rel.prop),
                            pn_stmt_value,
                        ]
                        for hop1_seq in hop1_seqs:
                            matches.append(DGPath(sequence=hop1_seq + hop2_seq))
                    else:
                        hop2_seq = [
                            DGPathEdge.p(rel.prop),
                            DGPathNodeStatement.from_FromWikidataLink(
                                middle_qnode.id, rel.prop, stmt_i
                            ),
                            DGPathEdge.p(rel.prop),
                        ]
                        for hop1_seq in hop1_seqs:
                            matches.append(DGPath(sequence=hop1_seq + hop2_seq))

                    for qual in rel.quals:
                        hop2_seq = [
                            DGPathEdge.p(rel.prop),
                            DGPathNodeStatement.from_FromWikidataLink(
                                middle_qnode.id, rel.prop, stmt_i
                            ),
                            DGPathEdge.p(qual),
                        ]
                        for hop1_seq in hop1_seqs:
                            matches.append(DGPath(sequence=hop1_seq + hop2_seq))

        return matches

    if options & BuildDGOption.NO_KG_INDEX:
        _path_object_search = _path_object_search_v1
    else:
        _path_object_search = _path_object_search_v2

    textparser = TextParser()
    funcs = {
        WikidataValueType.string.value: [LiteralMatcher.string_test_exact],
        WikidataValueType.globe_coordinate.value: [
            LiteralMatcher.globe_coordinate_test
        ],
        WikidataValueType.time.value: [LiteralMatcher.time_test],
        WikidataValueType.quantity.value: [LiteralMatcher.quantity_test],
        WikidataValueType.mono_lingual_text.value: [
            LiteralMatcher.mono_lingual_text_test_exact
        ],
        WikidataValueType.entity_id.value: [],  # [LiteralMatcher.entity_id_test_fuzzy]
    }

    # noinspection PyUnusedLocal,PyShadowingNames
    def _path_value_search(qnodes: Dict[str, QNode], source: QNode, value):
        matches = []
        for p, stmts in source.props.items():
            if p == "P31":
                # no need to search in the instanceOf property, as the ontology is removed from the databased as they are huge
                continue

            for stmt_i, stmt in enumerate(stmts):
                has_stmt_value = False
                if stmt.value.type in LiteralMatcher.literal_types:
                    stmt_value = stmt.value
                else:
                    assert stmt.value.type in LiteralMatcher.non_literal_types
                    if not stmt.value.is_qnode():
                        # lexical
                        continue
                    if stmt.value.as_entity_id() not in qnodes:
                        # this can happen due to some of the qnodes is in the link, but is missing in the KG
                        # this is very rare so we can employ some check to make sure this is not due to
                        # our wikidata subset
                        is_error_in_kg = any(
                            any(
                                _s.value.is_qnode()
                                and _s.value.as_entity_id() in qnodes
                                for _s in _stmts
                            )
                            for _p, _stmts in source.props.items()
                        ) or stmt.value.as_entity_id().startswith("L")
                        if not is_error_in_kg:
                            raise Exception(
                                f"Missing qnodes in your KG subset: {stmt.value.as_entity_id()}"
                            )
                        continue
                    stmt_value = qnodes[stmt.value.as_entity_id()]
                for fn in funcs[stmt.value.type]:
                    match, confidence = fn(stmt_value, value)
                    if match:
                        matches.append(
                            DGPath(
                                sequence=[
                                    DGPathEdge.p(p),
                                    DGPathNodeStatement.from_FromLiteralMatchingFunc(
                                        source.id,
                                        p,
                                        stmt_i,
                                        {"func": fn.__name__, "value": stmt.value},
                                        confidence,
                                    ),
                                    DGPathEdge.p(p),
                                ]
                            )
                        )
                        has_stmt_value = True

                for q, qvals in stmt.qualifiers.items():
                    for qval in qvals:
                        if qval.type in LiteralMatcher.literal_types:
                            qval_value = qval
                        else:
                            assert qval.type in LiteralMatcher.non_literal_types
                            # TODO: fix me handle L not correct
                            if qval.value["id"].startswith("L") or qval.value[
                                "id"
                            ].startswith("P"):
                                continue
                            if qval.as_entity_id() not in qnodes:
                                # this can happen due to some of the qnodes is in the link, but is missing in the KG
                                # this is very rare so we can employ some check to make sure this is not due to
                                # our wikidata subset
                                is_error_in_kg = any(
                                    any(
                                        _s.value.is_qnode()
                                        and _s.value.as_entity_id() in qnodes
                                        for _s in _stmts
                                    )
                                    for _p, _stmts in source.props.items()
                                ) or not qval.as_entity_id().startswith("L")
                                if not is_error_in_kg:
                                    raise Exception(
                                        f"Missing qnodes in your KG subset: {qval.as_entity_id()}"
                                    )
                                continue
                            qval_value = qnodes[qval.as_entity_id()]
                        for fn in funcs[qval.type]:
                            match, confidence = fn(qval_value, value)
                            if match:
                                if not has_stmt_value:
                                    if stmt.value.is_qnode():
                                        pn_stmt_value = DGPathNodeQNode(
                                            stmt.value.as_entity_id()
                                        )
                                    else:
                                        pn_stmt_value = DGPathNodeLiteralValue(
                                            stmt.value
                                        )
                                    matches.append(
                                        DGPath(
                                            sequence=[
                                                DGPathEdge.p(p),
                                                DGPathNodeStatement.from_FromWikidataLink(
                                                    source.id, p, stmt_i
                                                ),
                                                DGPathEdge.p(p),
                                                pn_stmt_value,
                                            ]
                                        )
                                    )
                                    has_stmt_value = True

                                matches.append(
                                    DGPath(
                                        sequence=[
                                            DGPathEdge.p(p),
                                            DGPathNodeStatement.from_FromLiteralMatchingFunc(
                                                source.id,
                                                p,
                                                stmt_i,
                                                {"func": fn.__name__, "value": qval},
                                                confidence,
                                            ),
                                            DGPathEdge.q(q),
                                        ]
                                    )
                                )
        return matches

    paths = []

    # no qnode in the source
    u_qnodes = (
        [qnodes[u.qnode_id]]
        if u.is_entity_value
        else [qnodes[qnode_id] for qnode_id in u.qnode_ids]
    )
    if len(u_qnodes) == 0:
        if bidirection:
            return kg_path_discovering(
                options,
                qnodes,
                kg_object_index,
                dg,
                v,
                u,
                max_n_hop=max_n_hop,
                bidirection=False,
            )
        return []

    v_qnodes = (
        [qnodes[v.qnode_id]]
        if v.is_entity_value
        else [qnodes[qnode_id] for qnode_id in v.qnode_ids]
    )
    v_value = v.context_span.get_text_span() if v.is_entity_value else v.value
    for n1 in u_qnodes:
        # count = timer.start("object discovery")
        for n2 in v_qnodes:
            for newpath in _path_object_search(
                qnodes, kg_object_index, n1, n2, max_n_hop=max_n_hop
            ):
                newpath.sequence.insert(0, DGPathExistingNode(u.id))
                if isinstance(newpath.sequence[-1], DGPathEdge):
                    newpath.sequence.append(DGPathExistingNode(v.id))
                paths.append(newpath)
        # count.stop()

        # count = timer.start("literal discovery")
        for newpath in _path_value_search(qnodes, n1, textparser.parse(v_value)):
            newpath.sequence.insert(0, DGPathExistingNode(u.id))
            if isinstance(newpath.sequence[-1], DGPathEdge):
                newpath.sequence.append(DGPathExistingNode(v.id))
            paths.append(newpath)
        # count.stop()

    if bidirection:
        paths += kg_path_discovering(
            options,
            qnodes,
            kg_object_index,
            dg,
            v,
            u,
            max_n_hop=max_n_hop,
            bidirection=False,
        )

    return paths


class KGInference:
    def __init__(
        self,
        dg: nx.MultiDiGraph,
        qnodes: Dict[str, QNode],
        wdprops: Dict[str, WDProperty],
    ):
        # mapping from qnode id, and property to a list of statement values (with the corresponding DG statement node if exist)
        # the reason we want to include all is that we want to know if we need to add new statement value or reuse existing value
        self.subkg: Dict[
            Tuple[str, str], List[Tuple[Optional[StatementNode], DataValue]]
        ] = {}
        self.qnodes = qnodes
        self.wdprops = wdprops
        self.dg = dg

        for sid, sdata in dg.nodes(data=True):
            s: StatementNode = sdata["data"]
            if not s.is_statement:
                continue

            dgsid = DGStatementID.parse_id(sid)
            self._set_stmt_node(s.qnode_id, dgsid.predicate, dgsid.statement_index, s)

    def infer_subproperty(self):
        """Infer new properties via sub-property of (inheritance)"""
        properties = set()
        qualifiers = set()

        for uid, udata in self.dg.nodes(data=True):
            u = udata["data"]
            if u.is_statement:
                continue

            for sid, us_edges in self.dg[u.id].items():
                # add all links to the list of properties
                properties.update(us_edges.keys())

                for vid, sv_edges in self.dg[sid].items():
                    qualifiers.update(
                        (v_eid for v_eid in sv_edges.keys() if v_eid not in us_edges)
                    )

        parent_props = self._build_parent_map(properties)
        parent_qualifiers = self._build_parent_map(qualifiers)

        # list of new properties and qualifiers that we will infer
        new_props = []
        new_qualifiers = []

        for sid, sdata in self.dg.nodes(data=True):
            s: StatementNode = sdata["data"]
            if not s.is_statement:
                continue

            # parents of the statement
            parents = set()
            prop = s.predicate
            for uid, _, eid, edata in self.dg.in_edges(sid, data=True, keys=True):
                parents.add(uid)
                assert prop == edata["data"].predicate

            # children that are properties of the statement
            prop_children = set()
            # children that are qualifiers of the statement
            qualifier2children = defaultdict(set)
            for _, vid, eid, edata in self.dg.out_edges(sid, data=True, keys=True):
                e: DGEdge = edata["data"]
                if e.predicate == prop:
                    prop_children.add(vid)
                else:
                    qualifier2children[e.predicate].add(vid)

            stmt_index = DGStatementID.parse_id(sid).statement_index
            if prop in parent_props:
                for parent_prop in parent_props[prop]:
                    for vid in prop_children:
                        for source_flow, flow_provenances in s.iter_source_flow(
                            EdgeFlowTarget(vid, prop)
                        ):
                            new_prop = InferredNewProp(
                                qnode_id=s.qnode_id,
                                new_prop=parent_prop,
                                value=self._get_stmt_value(
                                    s.qnode_id, prop, stmt_index
                                ),
                                source_id=source_flow.source_id,
                                target_id=vid,
                                qualifier_edges=[],
                                flow_provenances=flow_provenances,
                            )
                            new_props.append(new_prop)

            for q, children in qualifier2children.items():
                for pq in parent_qualifiers.get(q, []):
                    for vid in children:
                        for source_flow, provenance in s.iter_source_flow(
                            EdgeFlowTarget(vid, q)
                        ):
                            new_qualifiers.append(
                                InferredNewQualifier(
                                    statement_id=sid,
                                    new_qualifier=pq,
                                    source_id=source_flow.source_id,
                                    property=source_flow.edge_id,
                                    target_id=vid,
                                    flow_provenances=provenance,
                                )
                            )

        self.add_inference(new_props, new_qualifiers)
        return self

    def kg_transitive_inference(self):
        """Infer new relationship based on the transitive property: a -> b -> c => a -> c"""

        # find the list of transitive properties in the graph
        transitive_props = set()
        for uid, vid, eid, edata in self.dg.edges(data=True, keys=True):
            prop = self.wdprops[eid]
            # transitive class
            if "Q18647515" in prop.instanceof:
                transitive_props.add(eid)

        # now start from node u, we find if there is another v connect to u via a transitive property, and another p connect
        # to v with the same property, we don't need to keep the chain going as even if it's longer, we will eventually loop
        # through all item in the chain by looping through nodes in the graph

        chains = []
        for uid, udata in self.dg.nodes(data=True):
            u: DGNode = udata["data"]
            if u.is_statement:
                continue

            for sid, us_edges in self.dg[uid].items():
                stmt: StatementNode = self.dg.nodes[sid]["data"]
                for trans_prop in transitive_props:
                    if trans_prop not in us_edges:
                        continue
                    us_edge: DGEdge = us_edges[trans_prop]["data"]
                    for vid, sv_edges in self.dg[sid].items():
                        if trans_prop not in sv_edges:
                            continue
                        sv_edge: DGEdge = sv_edges[trans_prop]["data"]
                        if not stmt.is_same_flow(us_edge, sv_edge):
                            # don't allow infer new link across rows
                            continue

                        for s2id, vs2_edges in self.dg[vid].items():
                            if trans_prop not in vs2_edges:
                                continue
                            vs2_edge: DGEdge = vs2_edges[trans_prop]["data"]
                            stmt2: StatementNode = self.dg.nodes[s2id]["data"]
                            for v2id, s2v2_edges in self.dg[s2id].items():
                                if trans_prop not in s2v2_edges:
                                    continue
                                s2v2_edge: DGEdge = s2v2_edges[trans_prop]["data"]
                                if not stmt2.is_same_flow(vs2_edge, s2v2_edge):
                                    # don't allow infer new link across rows
                                    continue

                                # we now record the chain
                                chains.append((us_edge, sv_edge, vs2_edge, s2v2_edge))

        # make sure that there is no duplication in the chains
        assert len(chains) == len(
            {
                (
                    us_edge.source,
                    us_edge.predicate,
                    us_edge.target,
                    sv_edge.source,
                    sv_edge.predicate,
                    sv_edge.target,
                    vs2_edge.source,
                    vs2_edge.predicate,
                    vs2_edge.target,
                    s2v2_edge.source,
                    s2v2_edge.predicate,
                    s2v2_edge.target,
                )
                for us_edge, sv_edge, vs2_edge, s2v2_edge in chains
            }
        )

        # generate new property, but qualifiers cannot inherit via transitive inference
        new_props = []
        for us_edge, sv_edge, vs2_edge, s2v2_edge in chains:
            trans_prop = us_edge.predicate
            stmt: StatementNode = self.dg.nodes[us_edge.target]["data"]
            stmt2: StatementNode = self.dg.nodes[vs2_edge.target]["data"]

            prop_value = self._get_stmt_value(
                stmt2.qnode_id,
                trans_prop,
                DGStatementID.parse_id(stmt2.id).statement_index,
            )
            # prop_value = stmt2.qnode.props[trans_prop][DGStatementID.parse_id(stmt2.id).statement_index].value

            # calculating the provenance of the new transitive link. however, I don't know in case we have multiple
            # provenance since it depends on both first leg and second leg. for now, we put an assertion to handle
            # only the case where we have two legs has the same provenance, which make the new transitive link has
            # the same provenance too.
            first_leg_provenances = stmt.get_provenance_by_edge(us_edge, sv_edge)
            second_leg_provenances = stmt2.get_provenance_by_edge(vs2_edge, s2v2_edge)
            assert (
                len(first_leg_provenances) == 1
                and len(second_leg_provenances) == 1
                and first_leg_provenances[0] == second_leg_provenances[0]
            )

            provenances = first_leg_provenances
            new_props.append(
                InferredNewProp(
                    qnode_id=stmt.qnode_id,
                    new_prop=trans_prop,
                    value=prop_value,
                    source_id=us_edge.source,
                    target_id=s2v2_edge.target,
                    qualifier_edges=[],
                    flow_provenances=provenances,
                )
            )

        self.add_inference(new_props, [])
        return self

    def add_inference(
        self,
        new_props: List[InferredNewProp],
        new_qualifiers: List[InferredNewQualifier],
    ):
        """After we run inference, we got a list of new properties and new qualifiers that can update using this function"""
        new_nodes = []
        new_edges = []

        # here we enforce the constraint that there is no cross links
        # between nodes in different rows, this happen because transitive
        # inference generate links cross rows
        new_props = [
            new_prop
            for new_prop in new_props
            if DGEdge.can_link(
                self.dg.nodes[new_prop.source_id]["data"],
                self.dg.nodes[new_prop.target_id]["data"],
            )
        ]
        new_qualifiers = [
            new_qualifier
            for new_qualifier in new_qualifiers
            if DGEdge.can_link(
                self.dg.nodes[new_qualifier.source_id]["data"],
                self.dg.nodes[new_qualifier.target_id]["data"],
            )
        ]

        for new_prop in new_props:
            stmt_exist = None
            prop = new_prop.new_prop

            # search for existing statement in the KG
            self._track_property(new_prop.qnode_id, prop)
            for sprime, value in self.subkg[new_prop.qnode_id, prop]:
                if sprime is not None and new_prop.value == value:
                    stmt_exist = sprime
                    break

            # if the statement exist, re-use it
            if stmt_exist is not None:
                sprime = stmt_exist
            else:
                for stmt_index, (sprime, value) in enumerate(
                    self.subkg[new_prop.qnode_id, prop]
                ):
                    if new_prop.value == value:
                        stmt_id = DGStatementID(
                            new_prop.qnode_id, prop, stmt_index
                        ).get_id()
                        sprime = StatementNode(
                            stmt_id, new_prop.qnode_id, prop, is_in_kg=True
                        )
                        self._set_stmt_node(new_prop.qnode_id, prop, stmt_index, sprime)
                        new_nodes.append(sprime)
                        break
                else:
                    stmt_index = self.get_next_available_stmt_index(
                        new_prop.qnode_id, prop
                    )
                    stmt_id = DGStatementID(
                        new_prop.qnode_id, prop, stmt_index
                    ).get_id()
                    sprime = StatementNode(
                        stmt_id, new_prop.qnode_id, prop, is_in_kg=False
                    )
                    self._add_stmt_value(
                        new_prop.qnode_id, prop, stmt_index, sprime, new_prop.value
                    )
                new_nodes.append(sprime)

            if not self.dg.has_edge(new_prop.source_id, sprime.id, key=prop):
                new_edges.append(
                    DGEdge(
                        source=new_prop.source_id,
                        target=sprime.id,
                        predicate=prop,
                        is_qualifier=False,
                        is_inferred=True,
                    )
                )
            if not self.dg.has_edge(sprime.id, new_prop.target_id, key=prop):
                new_edges.append(
                    DGEdge(
                        source=sprime.id,
                        target=new_prop.target_id,
                        predicate=prop,
                        is_qualifier=False,
                        is_inferred=True,
                    )
                )

            sprime.track_provenance(
                EdgeFlowSource(new_prop.source_id, prop),
                EdgeFlowTarget(new_prop.target_id, prop),
                new_prop.flow_provenances,
            )

            # TODO: we haven't add to add add qualifiers, so we assert we don't have any. fix me!
            assert len(new_prop.qualifier_edges) == 0
            # for qual_edge in new_prop.qualifier_edges:
            #     new_edges.append(DGEdge(source=stmt_id, target=qual_edge.target, predicate=qual_edge.predicate, is_qualifier=True, paths=[], is_inferred=True))

        for new_qual in new_qualifiers:
            stmt: StatementNode = self.dg.nodes[new_qual.statement_id]["data"]
            new_edges.append(
                DGEdge(
                    source=new_qual.statement_id,
                    target=new_qual.target_id,
                    predicate=new_qual.new_qualifier,
                    is_qualifier=True,
                    is_inferred=True,
                )
            )
            stmt.track_provenance(
                EdgeFlowSource(new_qual.source_id, new_qual.property),
                EdgeFlowTarget(new_qual.target_id, new_qual.new_qualifier),
                new_qual.flow_provenances,
            )

        for node in new_nodes:
            self.dg.add_node(node.id, data=node)
        for edge in new_edges:
            self.dg.add_edge(edge.source, edge.target, key=edge.predicate, data=edge)

    def _track_property(self, qnode_id: str, prop: str):
        """Ensure that the subkg has values of the qnode's property"""
        if (qnode_id, prop) not in self.subkg:
            lst = []
            for stmt_i, stmt in enumerate(self.qnodes[qnode_id].props.get(prop, [])):
                lst.append((None, stmt.value))
            self.subkg[qnode_id, prop] = lst

    def _get_stmt_value(self, qnode_id: str, prop: str, stmt_index: int):
        self._track_property(qnode_id, prop)
        return self.subkg[qnode_id, prop][stmt_index][1]

    def _set_stmt_node(
        self, qnode_id: str, prop: str, stmt_index: int, stmt: StatementNode
    ):
        self._track_property(qnode_id, prop)
        assert (
            self.subkg[qnode_id, prop][stmt_index][0] is None
        ), "Cannot override existing value in the KG"
        self.subkg[qnode_id, prop][stmt_index] = (
            stmt,
            self.subkg[qnode_id, prop][stmt_index][1],
        )

    def _add_stmt_value(
        self,
        qnode_id: str,
        prop: str,
        stmt_index: int,
        stmt: StatementNode,
        value: DataValue,
    ):
        self._track_property(qnode_id, prop)
        assert stmt_index == len(self.subkg[qnode_id, prop]), "Can only add new value"
        self.subkg[qnode_id, prop].append((stmt, value))

    def get_next_available_stmt_index(self, qnode_id: str, prop: str):
        return len(self.subkg[qnode_id, prop])

    def _build_parent_map(self, props: Set[str]):
        """Build a map from a property to its parents in the same list"""
        parent_props: Dict[str, List[str]] = {}
        for p1 in props:
            parent_props[p1] = []
            for p2 in props:
                if p1 == p2:
                    continue
                if p2 in self.wdprops[p1].parents_closure:
                    parent_props[p1].append(p2)
            if len(parent_props[p1]) == 0:
                parent_props.pop(p1)
        return parent_props


class DGPruning:
    NxDGEdgeAttr = TypedDict("NxDGEdgeAttr", data=DGEdge)
    NxDGEdge = Tuple[str, str, str, NxDGEdgeAttr]

    def __init__(self, dg: nx.MultiDiGraph):
        self.dg = dg

    def prune_hidden_entities(self):
        """Prune redundant KG entities, which added to the graph via KG discovering and from the context.

        **Step 1:**
        Let:
        - n be an entity node in DG.
        - v is a node connected from n via a property: LEG2: n -> p -> s -> p' -> v, and s does not have other property/qualifier rather than p'

        We made the following heuristics:
        * If there is no other node connect to n, then n is a root node and is from the context. We should not
        prune this node, so just skip it.
        * For all another node ui \in U that connects to n via the path: LEG1: ui -> pi -> s' -> pi' -> n, if there is always a better
        path LEG* between ui and v, then we can remove the path LEG2. U contains nodes in cells or context, otherwise
        ui will be an entity to entity that won't be in the final model anyway.
        Note: LEG* is better than LEG2 when it's shorter, also from wikidata link or if not, it must have better match confidence (i.e., better provenance)

        **Step 2:**
        Let n' be an entity node in DG that do not link to other nodes (v doesn't exist).
        We have the following heuristics:
        * If there is no other node connect to it, this is a standable node and should be removed
        * For any node ui that connects to n via the path: LEG1: ui -> pi -> s' -> pi' -> n. If s' doesn't have other properties/qualifiers,
        then we can remove LEG1.

        **Step 3:**
        Finally, if a node is standable, we should remove it.
        """
        # step 1: prune the second leg paths
        legprime: Dict[Tuple[str, str], FlowProvenance] = {}
        rm_legs: List[Tuple[str, EdgeFlowSource, EdgeFlowTarget]] = []
        for nid, ndata in self.dg.nodes(data=True):
            n: EntityValueNode = ndata["data"]
            if not n.is_entity_value:
                continue

            if self.dg.in_degree(nid) == 0:
                # no other node connect to it
                continue

            # get list of grandpa ui (only cells or values in the context), along with their paths to node n.
            grandpa = set()
            for gpid in self.iter_grand_parents(nid):
                gp = self.dg.nodes[gpid]["data"]
                if gp.is_cell or gp.is_context:
                    grandpa.add(gpid)

            for _, sid, ns_eid in self.dg.out_edges(nid, keys=True):
                stmt: StatementNode = self.dg.nodes[sid]["data"]
                stmt_outedges = self.out_edges(sid)
                if len(stmt_outedges) > 1:
                    # this stmt has other qualifiers, so it's not what we are looking for
                    continue

                for sv_outedge in next(iter(stmt_outedges.values())):
                    v = self.dg.nodes[sv_outedge.target]["data"]
                    # got leg 2, now looks for all incoming
                    leg2 = (
                        EdgeFlowSource(nid, ns_eid),
                        EdgeFlowTarget(v.id, sv_outedge.predicate),
                    )
                    if not stmt.has_flow(*leg2):
                        continue
                    leg2_provenance = stmt.get_provenance(*leg2)

                    has_better_paths = True
                    for gpid in grandpa:
                        if (gpid, v.id) not in legprime:
                            paths = [
                                (
                                    self.dg.nodes[path[0][1]]["data"],
                                    EdgeFlowSource(path[0][0], path[0][2]),
                                    EdgeFlowTarget(path[1][1], path[1][2]),
                                )
                                for path in nx.all_simple_edge_paths(
                                    self.dg, gpid, v.id, cutoff=2
                                )
                            ]
                            provs = [
                                prov
                                for s, sf, tf in paths
                                if s.has_flow(sf, tf)
                                for prov in s.get_provenance(sf, tf)
                            ]
                            if len(provs) == 0:
                                legprime[gpid, v.id] = None
                            else:
                                legprime[gpid, v.id] = max(
                                    provs,
                                    key=cmp_to_key(
                                        self.specific_pruning_provenance_cmp
                                    ),
                                )
                        best_prov = legprime[gpid, v.id]
                        if (
                            best_prov is None
                            or max(
                                self.specific_pruning_provenance_cmp(
                                    best_prov, leg2_prov
                                )
                                for leg2_prov in leg2_provenance
                            )
                            < 0
                        ):
                            # no better path
                            has_better_paths = False
                            break

                    if has_better_paths:
                        rm_legs.append((sid, leg2[0], leg2[1]))

        # logger.info("#legs: {}", len(rm_legs))
        self.remove_flow(rm_legs)
        # logger.info("# 0-indegree: {}", sum(self.dg.in_degree(uid) == 0 for uid in self.dg.nodes))
        # logger.info("# 0-outdegree: {}", sum(self.dg.out_degree(uid) == 0 for uid in self.dg.nodes))
        # logger.info("# 0-standalone: {}",
        #             sum(self.dg.out_degree(uid) + self.dg.in_degree(uid) == 0 for uid in self.dg.nodes))

        # step 2: prune the first leg paths
        rm_legs: List[Tuple[str, EdgeFlowSource, EdgeFlowTarget]] = []
        for nid, ndata in self.dg.nodes(data=True):
            n: EntityValueNode = ndata["data"]
            if not n.is_entity_value or self.dg.out_degree(nid) > 0:
                continue

            for sid, _, sn_eid, sn_edata in self.dg.in_edges(nid, data=True, keys=True):
                stmt_outedges = self.out_edges(sid)
                if len(stmt_outedges) == 1:
                    # stmt does not have other property/qualifier
                    target_flow = EdgeFlowTarget(nid, sn_eid)
                    stmt: StatementNode = self.dg.nodes[sid]["data"]
                    for source_flow, _ in stmt.iter_source_flow(target_flow):
                        rm_legs.append((sid, source_flow, target_flow))

        # logger.info("#legs: {}", len(rm_legs))
        self.remove_flow(rm_legs)
        # logger.info("# 0-indegree: {}", sum(self.dg.in_degree(uid) == 0 for uid in self.dg.nodes))
        # logger.info("# 0-outdegree: {}", sum(self.dg.out_degree(uid) == 0 for uid in self.dg.nodes))
        # logger.info("# 0-standalone: {}",
        #             sum(self.dg.out_degree(uid) + self.dg.in_degree(uid) == 0 for uid in self.dg.nodes))
        self.prune_disconnected_nodes()

    def prune_disconnected_nodes(self):
        """This function prune out disconnected nodes that are:
        1. nodes without incoming edges and outgoing edges
        2. statement nodes with no incoming edges or no outgoing edges

        Returns
        -------
        """
        rm_nodes = set()
        for uid, udata in self.dg.nodes(data=True):
            u: DGNode = udata["data"]
            if u.is_entity_value:
                if self.dg.in_degree(uid) == 0 and self.dg.out_degree(uid) == 0:
                    rm_nodes.add(uid)
            elif u.is_statement:
                if self.dg.in_degree(uid) == 0 or self.dg.out_degree(uid) == 0:
                    rm_nodes.add(uid)
        for uid in rm_nodes:
            self.dg.remove_node(uid)

    def remove_flow(self, flows: List[Tuple[str, EdgeFlowSource, EdgeFlowTarget]]):
        for sid, source_flow, target_flow in flows:
            stmt: StatementNode = self.dg.nodes[sid]["data"]
            stmt.untrack_flow(source_flow, target_flow)
            if not stmt.has_source_flow(source_flow):
                self.dg.remove_edge(source_flow.source_id, sid, source_flow.edge_id)
            if not stmt.has_target_flow(target_flow):
                self.dg.remove_edge(sid, target_flow.target_id, target_flow.edge_id)

    def specific_pruning_provenance_cmp(
        self, prov0: FlowProvenance, prov1: FlowProvenance
    ):
        # compare provenance, this function only accept
        if prov0.gen_method == LinkGenMethod.FromWikidataLink:
            # always favour from wikidata link
            return 1
        if prov1.gen_method == LinkGenMethod.FromWikidataLink:
            return -1
        # assert prov0.gen_method == prov1.gen_method and prov0.gen_method_arg == prov1.gen_method_arg
        # do not need to check if the two gen method and args are equal, as even if we select the incorrect one
        # we only truncate when the other leg worst than it
        return prov0.prob - prov1.prob

    def iter_grand_parents(self, nid: str):
        for pid, _ in self.dg.in_edges(nid):
            for ppid, _ in self.dg.in_edges(pid):
                yield ppid

    def out_edges(self, uid: str) -> Dict[str, List[DGEdge]]:
        label2edges = defaultdict(list)
        for _, vid, eid, edata in self.dg.out_edges(uid, data=True, keys=True):
            label2edges[eid].append(edata["data"])
        return label2edges


def prune_data_graph_deprecated(dg: nx.MultiDiGraph):
    """This algorithm prunes:

    * Redundant hidden entities added to the graph from KG discovering. This can also include entities in the context.


    0. A second leg of a path between two nodes that goes through an entity node f1:
        (f0 -> p1 -> s1 -> p2 -> f1) (-> p3 -> s2 -> p3 -> f3). This removal is triggered when
        s2 doesn't have other qualifiers map to other nodes
    there is another better path between f1 and f3 (shorter, also from wikidatalink), s2 must not have any

    1. An entity that between all sources and targets that it connects exist a better path in the graph connecting them,
    and it is not used as a mean (should we consider it to be the only mean?) to connect the other columns/context as property & qualifiers (i.e., property value of statement). for example,
    an award is needed to connect the year column in the model. Note: A path is consider better when it is shorter,
    and also generated from WikidataLink rather than literal matching.
    2. Statements that belong to the same qnode and same property, have the same value and qualifiers matched in the data graph.
    The reason we may have multiple statements like this, is because a property may have multiple statements of the same values in
    KG, usually different with each other by the qualifiers

    Finally, after applying those pruning rules, if any dangling nodes will be removed too (e.g., a statement node that
    doesn't have any value, or an entity node that doesn't link to any node)

    Parameters
    ----------
    dg

    Returns
    -------
    """

    def iter_grand_parents(edges: List[DGPruning.NxDGEdge]):
        for f2, f3, eid, edge in edges:
            for f1, _f2 in dg.in_edges(f2):
                yield dg.nodes[f1]["data"]

    def iter_grand_children(edges: List[DGPruning.NxDGEdge]):
        """Iterate through children of children of a given node"""
        for f1, f2, eid, edge in edges:
            for f3, f23_edges in dg[f2].items():
                yield dg.nodes[f3]["data"]

    def iter_sibling_edges(edges: List[DGPruning.NxDGEdge]):
        """Iter through out sibling nodes"""
        for f1, f2, eid, edge in edges:
            for f2_sib, sib_edges in dg[f1].items():
                for sib_eid, sib_edata in sib_edges.items():
                    yield sib_edata["data"]

    def remove_nodes(node_ids: Set[str], remove_dangling_statements: bool = True):
        for nid in node_ids:
            node = dg.nodes[nid]["data"]
            if node.is_statement:
                # if it is statements, just remove it
                dg.remove_node(nid)
            else:
                # when delete a node, we need to remove its flow also
                for sid, _, eid in dg.in_edges(nid, keys=True):
                    s: StatementNode = dg.nodes[sid]["data"]
                    s.untrack_target_flow(EdgeFlowTarget(nid, eid))
                for _, sid, eid in dg.out_edges(nid, keys=True):
                    s: StatementNode = dg.nodes[sid]["data"]
                    s.untrack_source_flow(EdgeFlowSource(nid, eid))
                dg.remove_node(nid)

        if not remove_dangling_statements or len(node_ids) == 0:
            return

        # remove dangling statements
        dangling_statements = set()
        for nid in dg.nodes:
            node = dg.nodes[nid]["data"]
            if node.is_statement and (
                dg.out_degree(nid) == 0 or dg.in_degree(nid) == 0
            ):
                dangling_statements.add(nid)
        for nid in dangling_statements:
            dg.remove_node(nid)

    def extract_stmt_signature(stmt_id: str):
        # to compare if two statements are equal, we only need to compare if the property & qualifiers values are
        # the same. We define that the value is the same by their node id in the graph (so that cell need to be the exact
        # cell, while literal or entity will be the value)
        signature = set()
        for uid, edges in dg[stmt_id].items():
            for eid in edges:
                signature.add(f"({eid}-{uid})")
        signature = sorted(list(signature))
        return "-".join(signature)

    has_direct_kglink_paths = {}

    rm_node_ids = set()
    # remove entities first.
    for uid, udata in dg.nodes(data=True):
        u: EntityValueNode = udata["data"]
        if not u.is_entity_value:
            continue

        # find entity node that only connected by two distinct qnodes
        inedges: List[DGPruning.NxDGEdge] = list(dg.in_edges(uid, data=True, keys=True))
        outedges: List[DGPruning.NxDGEdge] = list(
            dg.out_edges(uid, data=True, keys=True)
        )

        if len(inedges) == 0:
            # this happen when this entity is needed to links between columns or column and context as property & qualifiers
            # we can't prune this node, but we want to make sure this is actually the case
            assert (
                len(
                    {
                        v.column if v.is_cell else v.id
                        for v in iter_grand_children(outedges)
                        if v.is_cell
                        or ((v.is_literal_value or v.is_entity_value) and v.is_context)
                    }
                )
                >= 2
            ), "Should have at least two columns or one column one context"
            continue

        if len(outedges) == 0:
            # this happen when the real value of statement is missing, we can't prune this node
            # but we want to make sure there is actually another qualifiers are there
            assert (
                len({sib_edge.predicate for sib_edge in iter_sibling_edges(inedges)})
                >= 2
            ), "The statement should have at least two distinct predicates"
            continue

        # check if it is used to connect other columns via property & qualifiers

        # check if this entity only connect two qnodes and no literal values
        source_qnodes = {dg.nodes[sid]["data"].qnode_id for _1, sid, _2, _3 in inedges}
        if len(source_qnodes) != 1:
            continue
        # TODO: in order to get correct target qnodes, we need to examine the flow of the statements, the current finding
        # target qnodes will not work as expected when the target cell have more than one qnode
        target_qnodes = set()
        for v in iter_grand_children(outedges):
            if v.is_cell:
                target_qnodes = target_qnodes.union(v.qnode_ids)
            elif v.is_entity_value:
                target_qnodes.add(v.qnode_id)
            else:
                assert v.is_literal_value
                target_qnodes = set()
                break
        if len(target_qnodes) != 1:
            continue

        # now they connect only two distinct qnodes and no literal, so we know this path is through wikidata links
        # find out if all grand parents and grand children have better paths
        # first retrieve the list of grant parent and grant child we need to test.
        # instead of finding pairs of parents and children through flows, we do it here..?
        grant_parents = list(iter_grand_parents(inedges))
        grand_children = list(iter_grand_children(outedges))

        grand_child_cells = defaultdict(list)
        grand_child_ents = []
        for gc in grand_children:
            if gc.is_cell:
                grand_child_cells[gc.row].append(gc)
            else:
                assert gc.is_entity_value
                grand_child_ents.append(gc)

        all_paths_are_better = True
        for gp in grant_parents:
            if gp.is_cell:
                corresponding_gcs = chain(grand_child_cells[gp.row], grand_child_ents)
            else:
                corresponding_gcs = grand_child_ents
            for gc in corresponding_gcs:
                # identify if there is other better path is shorter and directly connect the two nodes
                if (gp.id, gc.id) not in has_direct_kglink_paths:
                    has_better_path = False
                    has_better_path = any(
                        (
                            any(
                                (
                                    provenance.gen_method
                                    == LinkGenMethod.FromWikidataLink
                                    for provenance in dg.nodes[path[0][1]][
                                        "data"
                                    ].get_provenance(
                                        EdgeFlowSource(path[0][0], path[0][2]),
                                        EdgeFlowTarget(path[1][1], path[1][2]),
                                    )
                                )
                            )
                            for path in nx.all_simple_edge_paths(
                                dg, gp.id, gc.id, cutoff=2
                            )
                        )
                    )
                    has_direct_kglink_paths[gp.id, gc.id] = has_better_path
                else:
                    has_better_path = has_direct_kglink_paths[gp.id, gc.id]
                if not has_better_path:
                    all_paths_are_better = False
                    break
            if not all_paths_are_better:
                break

        if all_paths_are_better:
            rm_node_ids.add(uid)

    remove_nodes(rm_node_ids)

    rm_node_ids = set()
    # remove duplicated statements
    for uid, udata in dg.nodes(data=True):
        u: Union[CellNode, EntityValueNode] = udata["data"]
        if not (u.is_cell or u.is_entity_value):
            continue

        # we are looking at statements of the same edges
        prop2stmts = defaultdict(list)
        for sid, us_edges in dg[uid].items():
            for eid, edata in us_edges.items():
                prop2stmts[eid].append(sid)

        for stmts in prop2stmts.values():
            # remove duplicated statements
            if len(stmts) > 1:
                stmt_sig = {}
                for si, sid in enumerate(stmts):
                    sig = extract_stmt_signature(sid)
                    if sig in stmt_sig:
                        continue
                    stmt_sig[sig] = si
                if len(stmt_sig) < len(stmts):
                    for si in set(range(len(stmts))).difference(stmt_sig.values()):
                        rm_node_ids.add(stmts[si])

    remove_nodes(rm_node_ids)


def build_data_graph(
    table: LinkedTable,
    qnodes: Dict[str, QNode],
    wdprops: Dict[str, WDProperty],
    kg_object_index: KGObjectIndex,
    ignore_columns: Set[int] = None,
    max_n_hop: int = 2,
    options: BuildDGOption = BuildDGOption.DISABLE_SAME_ENT_SEARCH
    | BuildDGOption.PRUNING_REDUNDANT_ENT,
    verbose: bool = False,
):
    """Build data graph from a linked table"""
    dg = nx.MultiDiGraph()
    context_node_id = None

    for ci, col in enumerate(table.table.columns):
        for ri, val in enumerate(col.values):
            # =====================================================
            # NOTE: old code
            # cell_qnodes = set()
            # cell_qnode_spans = {}
            # for link in table.links[ri][ci]:
            #     if link.entity_id is not None:
            #         cell_qnodes.add(link.entity_id)
            #         if link.entity_id not in cell_qnode_spans:
            #             cell_qnode_spans[link.entity_id] = []
            #         cell_qnode_spans[link.entity_id].append(Span(link.start, link.end))
            # assert all(len(spans) == len(set(spans)) for spans in cell_qnode_spans.values())
            # TODO: new code, doesn't handle the qnode_spans correctly
            cell_qnodes = {
                candidate.entity_id
                for link in table.links[ri][ci]
                for candidate in link.candidates
            }
            cell_qnode_spans = {}
            for link in table.links[ri][ci]:
                if link.entity_id is not None or len(link.candidates) > 0:
                    if len(link.candidates) > 0:
                        tmpid = link.candidates[0].entity_id
                    else:
                        tmpid = link.entity_id
                    if tmpid not in cell_qnode_spans:
                        cell_qnode_spans[tmpid] = []
                    cell_qnode_spans[tmpid].append(Span(link.start, link.end))
            assert all(
                len(spans) == len(set(spans)) for spans in cell_qnode_spans.values()
            )
            # =====================================================

            node = CellNode(
                id=f"{ri}-{ci}",
                value=val,
                column=ci,
                row=ri,
                qnode_ids=list(cell_qnodes),
                qnodes_span=cell_qnode_spans,
            )
            dg.add_node(node.id, data=node)

    if table.context.page_entity_id is not None:
        context_node_id = DGPathNodeQNode(table.context.page_entity_id).get_id()
        node = EntityValueNode(
            id=context_node_id,
            qnode_id=table.context.page_entity_id,
            context_span=ContextSpan(
                text=table.context.page_title,
                span=Span(0, len(table.context.page_title)),
            ),
        )
        dg.add_node(node.id, data=node)

    # find all paths
    n_rows = len(table.table.columns[0].values)
    kg_path_discovering_tasks = []
    new_paths = []
    for ri in range(n_rows):
        for ci, col in enumerate(table.table.columns):
            if ignore_columns is not None and ci in ignore_columns:
                continue

            u = dg.nodes[f"{ri}-{ci}"]["data"]
            for cj in range(ci + 1, len(table.table.columns)):
                if ignore_columns is not None and cj in ignore_columns:
                    continue

                v = dg.nodes[f"{ri}-{cj}"]["data"]
                kg_path_discovering_tasks.append((u, v))
            if context_node_id in dg:
                kg_path_discovering_tasks.append((u, dg.nodes[context_node_id]["data"]))

    for u, v in (
        tqdm(kg_path_discovering_tasks, desc="KG searching")
        if verbose
        else kg_path_discovering_tasks
    ):
        new_paths += kg_path_discovering(
            options, qnodes, kg_object_index, dg, u, v, max_n_hop=max_n_hop
        )

    # add paths to the graph
    for path in new_paths:
        curr_nodeid = path.sequence[0].id
        tmp_path = [curr_nodeid]
        for i in range(1, len(path.sequence), 2):
            prop: DGPathEdge = path.sequence[i]
            value: DGPathNode = path.sequence[i + 1]

            if isinstance(value, DGPathExistingNode):
                nodeid = value.id
            else:
                nodeid = value.get_id()
                if nodeid not in dg.nodes:
                    if isinstance(value, DGPathNodeStatement):
                        dg.add_node(
                            nodeid,
                            data=StatementNode(
                                nodeid,
                                qnode_id=value.qnode_id,
                                predicate=value.predicate,
                                is_in_kg=True,
                            ),
                        )
                    elif isinstance(value, DGPathNodeQNode):
                        dg.add_node(
                            nodeid,
                            data=EntityValueNode(
                                nodeid, qnode_id=value.qnode_id, context_span=None
                            ),
                        )
                    else:
                        dg.add_node(
                            nodeid,
                            data=LiteralValueNode(
                                nodeid, value=value.value, context_span=None
                            ),
                        )

            if dg.has_edge(curr_nodeid, nodeid, prop.value):
                edge: DGEdge = dg.edges[curr_nodeid, nodeid, prop.value]["data"]
            else:
                edge = DGEdge(
                    source=curr_nodeid,
                    target=nodeid,
                    predicate=prop.value,
                    is_qualifier=prop.is_qualifier,
                )
                dg.add_edge(curr_nodeid, nodeid, key=prop.value, data=edge)

            tmp_path.append(edge)
            tmp_path.append(nodeid)
            curr_nodeid = nodeid

        # path format: u - e - (s) - e - v - e - (s2) - e - v2
        for i in range(2, len(tmp_path), 4):
            u_node_id = tmp_path[i - 2]
            u_edge = tmp_path[i - 1]
            snode: StatementNode = dg.nodes[tmp_path[i]]["data"]
            v_edge = tmp_path[i + 1]
            v_node_id = tmp_path[i + 2]

            edge_source = EdgeFlowSource(u_node_id, u_edge.predicate)
            edge_target = EdgeFlowTarget(v_node_id, v_edge.predicate)
            snode.track_provenance(
                edge_source, edge_target, [path.sequence[i].provenance]
            )

    # DEBUG code
    # for uid, udata in dg.nodes(data=True):
    #     u = udata['data']
    #     if not u.is_cell:
    #         continue
    #     if u.column == 0:
    #         for _, sid, eid, edata in dg.out_edges(uid, data=True, keys=True):
    #             for _1, vid, e2, e2data in dg.out_edges(sid, data=True, keys=True):
    #                 v = dg.nodes[vid]['data']
    #                 if v.is_cell and v.column == 1:
    #                     print(uid, sid, vid, eid, e2)

    KGInference(dg, qnodes, wdprops).infer_subproperty().kg_transitive_inference()

    # pruning unnecessary paths
    if options & BuildDGOption.PRUNING_REDUNDANT_ENT:
        DGPruning(dg).prune_hidden_entities()

    M.log("grams", data_graph=dg)
    return dg
