"""
Lifts Data nodes to top-level GNF Logicals with fresh per-column predicates.

This pass transforms Data nodes that appear nested within Logicals into a set of
top-level GNF (Graph Normal Form) Logicals — one per column plus one for
the row ID.  Each GNF Logical contains the annotated Data node and a derive
statement for the corresponding 1- or 2-arity column predicate.  The original
Data node is replaced with a collection of Lookups to those column predicates.

Example transformation:

Before:
  root: Logical
    ...

    Logical
      Data{df}(row_id, col_1, col_2, col_3)
      predicate(col_1, col_3)

After:
  root: Logical
    Logical @lifted_data
      Data{df}(row_id, col_1, col_2, col_3)
      Logical @gnf_data("__ROW_ID__")
        derive _data_1___row_id__(row_id)
      Logical @gnf_data("col_1")
        derive _data_1_col_1(row_id, col_1)
      Logical @gnf_data("col_2")
        derive _data_1_col_2(row_id, col_2)
      Logical @gnf_data("col_3")
        derive _data_1_col_3(row_id, col_3)
    ...
    Logical
      Logical
        _data_1___row_id__(row_id)
        _data_1_col_1(row_id, col_1)
        _data_1_col_3(row_id, col_3)
      predicate(col_1, col_3)


Temporarily, for backward compatibility, we transform single column Data nodes with less
than 5 values (rows) into a disjunction of equality predicates.

Before:
  root:
    Logical
      ...
      Logical
        Data{[1, 2]}(x) <- small df, less than 5 values
        predicate(x)

After:
  root:
    Logical
      ...
      Logical
        Union
          Logical
            x = 1
          Logical
            x = 2
        predicate(x)

"""

from relationalai.semantics.metamodel.compiler import Pass
from relationalai.semantics.metamodel import ir, builtins as rel_builtins, factory as f, visitor
from relationalai.semantics.lqp.builtins import lifted_data_annotation, gnf_data_annotation, ROW_ID_COLUMN
import hashlib
import pandas as pd


class LiftData(Pass):
    def __init__(self):
        # Persistent counter across compilations so that data node names are globally unique
        # (the same LiftData instance is reused for model and query compilations).
        self.lifted_data_node_count = 0

    def rewrite(self, model: ir.Model, options: dict = {}) -> ir.Model:  # noqa: ARG002
        lifter = DataLifter(lifted_data_node_count=self.lifted_data_node_count)
        result = lifter.walk(model)
        self.lifted_data_node_count = lifter.lifted_data_node_count
        return result


class DataLifter(visitor.Rewriter):
    _EXPAND_MAX_ROWS = 5

    def __init__(self, lifted_data_node_count: int = 0):
        # GNF path state
        self.gnf_lifted_data: list[ir.Logical] = []   # one lifted data Logical (annotated @lifted_data) per unique Data
        self.gnf_relations: list[ir.Relation] = []   # all derived relations

        # Cache: Data node content hash -> [(ROW_ID_COLUMN, _data_k___row_id__), ("col_1", _data_k_col_1), ...]
        self.gnf_cache: dict[str, list[tuple[str, ir.Relation]]] = {}
        self.lifted_data_node_count = lifted_data_node_count

        super().__init__()

    def handle_model(self, model: ir.Model, parent: None) -> ir.Model:
        model = super().handle_model(model, parent)

        if not self.gnf_lifted_data:
            return model

        assert isinstance(model.root, ir.Logical)
        root = model.root
        return ir.Model(
            model.engines,
            model.relations | self.gnf_relations,
            model.types,
            ir.Logical(
                root.engine,
                root.hoisted,
                tuple(self.gnf_lifted_data) + root.body,
                root.annotations,
            ),
            model.annotations,
        )

    def handle_data(self, node: ir.Data, parent: ir.Node) -> ir.Logical:
        """Replace a Data node with a Logical containing Lookups (GNF) or a Union (backward compat)."""
        if self.should_expand(node):
            return self.expand_disjunction(node)
        else:
            return self.gnf_wrapper(node)

    # ------------------------------------------------------------------ #
    # GNF path                                                             #
    # ------------------------------------------------------------------ #

    def gnf_wrapper(self, node: ir.Data) -> ir.Logical:
        """Return a wrapper Logical whose body is the Lookup sequence for node."""
        return f.logical(self.gnf_lookups(node), list(node.vars))

    def gnf_lookups(self, node: ir.Data) -> list[ir.Task]:
        """Return [Lookup(row_id_rel, [row_id]), Lookup(col_rel, [row_id, col]), …]."""
        row_id_var = node.vars[0]
        signature = self.get_data_gnf_signature(node)
        # First entry is the 1-arity __ROW_ID__ predicate
        lookups: list[ir.Task] = [f.lookup(signature[0][1], [row_id_var])]
        # Remaining entries are 2-arity column predicates: (row_id, col_var)
        col_vars = [node.vars[i] for i in range(1, len(node.vars))]
        for (_, rel), var in zip(signature[1:], col_vars):
            lookups.append(f.lookup(rel, [row_id_var, var]))
        return lookups

    def get_data_gnf_signature(self, node: ir.Data) -> list[tuple[str, ir.Relation]]:
        """Return cached (name, relation) pairs for node, lifting it on first call.

        Returns:
            A list where index 0 is (ROW_ID_COLUMN, row_id_relation) — 1-arity,
            and indices 1+ are (col_name, col_relation) — 2-arity each.
        """
        cache_key = self.data_content_hash(node)
        if cache_key not in self.gnf_cache:
            self.gnf_cache[cache_key] = self.lift_data_gnf(node)
        return self.gnf_cache[cache_key]

    def lift_data_gnf(self, node: ir.Data) -> list[tuple[str, ir.Relation]]:
        """Lift a Data node: build per-column relations, register the lifted data Logical."""
        self.lifted_data_node_count += 1
        signature, gnf_children = self.mk_gnf_signature(node, self.lifted_data_node_count)
        self.gnf_lifted_data.append(f.logical(
            [node] + gnf_children,
            annos=[lifted_data_annotation()],
        ))
        return signature

    def mk_gnf_signature(
        self, node: ir.Data, n: int
    ) -> tuple[list[tuple[str, ir.Relation]], list[ir.Task]]:
        """Build per-column relations and @gnf_data derive nodes for node.

        Returns:
            (signature, gnf_children) where signature is the (name, relation) list
            and gnf_children are the @gnf_data derive nodes.
        """
        result: list[tuple[str, ir.Relation]] = []
        gnf_children: list[ir.Task] = []
        row_id_var = node.vars[0]

        # 1-arity row-ID predicate
        row_id_rel = f.relation(
            f"_data_{n}_{ROW_ID_COLUMN}",
            [f.field(row_id_var.name, row_id_var.type)],
        )
        self.gnf_relations.append(row_id_rel)
        result.append((ROW_ID_COLUMN, row_id_rel))
        gnf_children.append(f.logical(
            [f.derive(row_id_rel, [row_id_var])],
            annos=[gnf_data_annotation(ROW_ID_COLUMN)],
        ))

        # 2-arity per-column predicates
        col_vars = [node.vars[i] for i in range(1, len(node.vars))]
        for col_name, var in zip(node.data.columns, col_vars):
            col_rel = f.relation(
                f"_data_{n}_{col_name}",
                [
                    f.field(row_id_var.name, row_id_var.type),
                    f.field(var.name, var.type),
                ],
            )
            self.gnf_relations.append(col_rel)
            result.append((col_name, col_rel))
            gnf_children.append(f.logical(
                [f.derive(col_rel, [row_id_var, var])],
                annos=[gnf_data_annotation(col_name)],
            ))

        return result, gnf_children

    # ------------------------------------------------------------------ #
    # Backward compatibility path (single-column, < 5 rows)              #
    # These two methods form a self-contained code path that can be      #
    # deleted once the GNF path is fully adopted.                        #
    # ------------------------------------------------------------------ #

    @staticmethod
    def should_expand(node: ir.Data) -> bool:
        """True for single-column (excluding row_id) Data nodes with fewer than _EXPAND_MAX_ROWS rows."""
        return len(node.data.columns) == 1 and 0 < len(node.data) < DataLifter._EXPAND_MAX_ROWS

    @staticmethod
    def expand_disjunction(node: ir.Data) -> ir.Logical:
        """Expand a small single-column Data into a wrapper Logical containing a Union.

        The Union has one branch per row; each branch equates the row-id var and the
        column var to their literal values from the DataFrame.
        """
        row_id_var = node.vars[0]
        col_var = node.vars[1]
        branches: list[ir.Task] = []
        for row in node:  # itertuples(index=True): Pandas(Index=idx, col=val)
            idx = row[0]
            col_val = row[1]
            branch = f.logical(
                [
                    f.lookup(rel_builtins.eq, [row_id_var, f.literal(DataLifter.to_python(idx), type=row_id_var.type)]),
                    f.lookup(rel_builtins.eq, [col_var, f.literal(DataLifter.to_python(col_val), type=col_var.type)]),
                ],
                hoisted=list(node.vars),
            )
            branches.append(branch)
        union = f.union(branches, hoisted=list(node.vars))
        return f.logical([union])

    @staticmethod
    def to_python(val):
        """Convert a numpy scalar to a plain Python value."""
        return val.item() if hasattr(val, "item") else val

    # ------------------------------------------------------------------ #
    # Shared helpers                                                       #
    # ------------------------------------------------------------------ #

    @staticmethod
    def data_content_hash(node: ir.Data) -> str:
        """Content-based cache key incorporating column names, var types, and row data."""
        col_names = ",".join(str(c) for c in node.data.columns)
        col_types = ",".join(str(v.type) for v in node.vars)
        row_bytes = bytes(pd.util.hash_pandas_object(node.data, index=True).values)
        return hashlib.sha256(f"{col_names}|{col_types}|".encode() + row_bytes).hexdigest()
