from ..metamodel import ir, builtins, helpers, types
from ..metamodel.visitor import collect_by_type
from ..metamodel.util import FrozenOrderedSet, OrderedSet
from ..metamodel.compiler import group_tasks
from . import lqp_bundle
from . import utils, types as lqp_types
from .primitives import lqp_avg_op, lqp_operator, build_primitive
from .pragmas import pragma_to_lqp_name
from .types import meta_type_to_lqp
from .constructors import (
    mk_abstraction, mk_and, mk_exists, mk_or, mk_pragma, mk_primitive,
    mk_specialized_value, mk_type, mk_value, mk_attribute,
)
from .algorithms import (
    is_script, is_algorithm_script, is_algorithm_logical,
    is_while_loop, is_while_script, construct_monoid, is_empty_instruction
)
from .builtins import (
    has_global_annotation, get_upsert_annotation, get_monoid_annotation,
    get_monus_annotation, has_assign_annotation, get_arity, supported_lqp_annotations,
    is_lifted_data, has_gnf_data_annotation, get_gnf_column_name, ROW_ID_COLUMN,
)
from .utils import (
    TranslationCtx, ExportDescriptor, ExportInfo, gen_unique_var, gen_rel_id, extract_export_fqn
)
from .validators import assert_valid_input
from .rewrite.functional_dependencies import (
    normalized_fds, contains_only_declarable_constraints
)
from decimal import Decimal as PyDecimal
from datetime import datetime, date, timezone
from typing import Sequence, Tuple, cast, Union, Optional
from warnings import warn
import io
import pandas as pd
import re
import uuid
from collections import defaultdict

# Main access point for translating metamodel to lqp. Converts the model IR to an LQP epoch.
def to_lqp(model: ir.Model, fragment_name: bytes, ctx: TranslationCtx) -> tuple[Optional[ExportInfo], lqp_bundle.ir.Epoch]:
    assert_valid_input(model)
    decls: list[lqp_bundle.ir.Declaration] = []
    reads: list[lqp_bundle.ir.Read] = []

    # LQP only accepts logical tasks
    # These are asserted at init time
    root = cast(ir.Logical, model.root)
    for subtask in root.body:
        assert isinstance(subtask, ir.Logical)
        new_decls = _translate_to_decls(ctx, subtask)
        decls.extend(new_decls)

    # Add pyrel error attributes to reads.
    for err_id in _extract_pyrel_error_ids(ctx, model):
        ctx.output_ids.append(err_id)

    reads.extend(_get_output_reads(ctx.output_ids))

    # export_info is None if no exports, or a list of (fqn, filename, col_types) tuples
    export_info = None
    if len(ctx.export_descriptors) > 0:
        # Group descriptors by export FQN to create separate reads per export table
        exports_by_fqn = defaultdict(list)
        for descriptor in ctx.export_descriptors:
            exports_by_fqn[descriptor.export_fqn].append(descriptor)

        # Create a read for each export table and collect export info for each
        export_info = []
        for export_fqn, descriptors in exports_by_fqn.items():
            export_filename, col_types, export_reads = _get_export_reads(descriptors)
            reads.extend(export_reads)
            export_info.append((export_fqn, export_filename, col_types))

    debug_info = lqp_bundle.ir.DebugInfo(id_to_orig_name=ctx.rel_id_to_orig_name, meta=None)
    fragment_id = lqp_bundle.ir.FragmentId(id=fragment_name, meta=None)
    fragment = lqp_bundle.ir.Fragment(id=fragment_id, declarations=decls, meta=None, debug_info=debug_info)
    define_op = lqp_bundle.ir.Define(fragment=fragment, meta=None)

    txn = lqp_bundle.ir.Epoch(
        reads=reads,
        writes=[lqp_bundle.ir.Write(write_type=define_op, meta=None)],
        meta=None
    )

    return (export_info, txn)

def _effect_bindings(effect: Union[ir.Output, ir.Update]) -> list[ir.Value]:
    if isinstance(effect, ir.Output):
        # Unions may not return anything. The generated IR contains a None value when this
        # happens. We ignore it here.
        return [v for v in helpers.output_values(effect.aliases) if v]
    else:
        return list(effect.args)

def _get_output_reads(output_ids: list[tuple[lqp_bundle.ir.RelationId, str]]) -> list[lqp_bundle.ir.Read]:
    reads = []
    for (rel_id, name) in output_ids:
        assert isinstance(rel_id, lqp_bundle.ir.RelationId)
        output = lqp_bundle.ir.Output(name=name, relation_id=rel_id, meta=None)
        reads.append(lqp_bundle.ir.Read(read_type=output, meta=None))
    return reads

def _get_export_reads(descriptors: list[ExportDescriptor]) -> tuple[str, list, list[lqp_bundle.ir.Read]]:
    reads = []
    csv_columns = []
    col_info = []
    for descriptor in sorted(descriptors, key=lambda x: x.column_number):
        csv_columns.append(lqp_bundle.ir.ExportCSVColumn(
            column_name=descriptor.column_name,
            column_data=descriptor.relation_id,
            meta=None,
        ))
        col_info.append((descriptor.column_name, descriptor.column_type))

    # Generate a random name for the internal export path
    export_filename = "export_" + str(uuid.uuid4()).replace("-", "_")

    # Note that the engine will append the transaction id to the export path for access
    # control reasons. So the actual final path will be
    # `.../export/{export_filename}/data_{txn_id}`
    export_path = f"snowflake://APP_STATE.RAI_INTERNAL_STAGE/export/{export_filename}/data"
    export_csv_config = lqp_bundle.ir.ExportCSVConfig(
        path=export_path,
        data_columns=csv_columns,
        compression="gzip",
        partition_size=200,
        syntax_escapechar='"',  # To follow Snowflake's expected format
        meta=None,
    )
    reads.append(lqp_bundle.ir.Read(read_type=lqp_bundle.ir.Export(config=export_csv_config, meta=None), meta=None))
    return (export_filename, col_info, reads)

def _translate_to_decls(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp_bundle.ir.Declaration]:
    if is_lifted_data(rule):
        return _translate_lifted_data(ctx, rule)
    elif contains_only_declarable_constraints(rule):
        return _translate_to_constraint_decls(ctx, rule)
    elif is_algorithm_logical(rule):
        return _translate_algorithms(ctx, rule)
    else:
        return _translate_to_standard_decl(ctx, rule)

# Translate a lifted data Logical (annotated @lifted_data) into a CSVData declaration and per-column Defs.
# From
#     Logical @lifted_data
#         Data{df}(row_id, col_1, col_2, ...)
#         Logical @gnf_data("__ROW_ID__")
#             derive _data_k___row_id__(row_id)
#         Logical @gnf_data("col_1")
#             derive _data_k_col_1(row_id, col_1)
#         Logical @gnf_data("col_2")
#             derive _data_k_col_2(row_id, col_2)
#         Logical @gnf_data("col_3")
#             derive _data_k_col_3(row_id, col_3)
# To
#     (csv-data
#         (csv-locator :inline "<col_1>,<col_2>,<col_3>\n<row1>...\n<row2>...\n")
#         (csv-config :header-row 1 :delimiter "," :newline "\n" :quotechar "\"" :escapechar "\"" ...)
#         (csv-column "col_1" _data_k_col_1_col_col_1 (INT <col_1_type>))
#         (csv-column "col_2" _data_k_col_2_col_col_2 (INT <col_2_type>))
#         (csv-column "col_3" _data_k_col_3_col_col_3 (INT <col_3_type>)))
#     (def _data_k___row_id__ (row_id::INT128)
#         (exists (_row_id::INT, _v::<col_1_type>) (and (cast _row_id row_id) (atom _data_k_col_1_col_col_1 _row_id _v))))
#     (def _data_k_col_1 (row_id::INT128, col_1::<col_1_type>)
#         (exists (_row_id::INT) (and (cast _row_id row_id) (atom _data_k_col_1_col_col_1 _row_id col_1))))
#     (def _data_k_col_2 (row_id::INT128, col_2::<col_2_type>)
#         (exists (_row_id::INT) (and (cast _row_id row_id) (atom _data_k_col_2_col_col_2 _row_id col_2))))
#     (def _data_k_col_3 (row_id::INT128, col_3::<col_3_type>)
#         (exists (_row_id::INT) (and (cast _row_id row_id) (atom _data_k_col_3_col_col_3 _row_id col_3))))
#
def _translate_lifted_data(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp_bundle.ir.Declaration]:
    # Extract the Data node and @gnf_data children.
    data_node = None
    gnf_children: list[ir.Logical] = []
    for task in rule.body:
        if isinstance(task, ir.Data):
            assert data_node is None, "Expected exactly one Data node in lifted data Logical"
            data_node = task
        elif isinstance(task, ir.Logical) and has_gnf_data_annotation(task):
            gnf_children.append(task)
        else:
            raise ValueError(f"Unexpected task in lifted data Logical: {type(task)}")
    assert data_node is not None, "Missing Data node in lifted data Logical"

    # Separate @gnf_data children into the __ROW_ID__ entry and column entries.
    row_id_gnf = None
    col_gnfs: list[tuple[str, ir.Logical]] = []
    for gnf in gnf_children:
        col_name = get_gnf_column_name(gnf)
        if col_name == ROW_ID_COLUMN:
            row_id_gnf = gnf
        else:
            col_gnfs.append((col_name, gnf))
    assert row_id_gnf is not None, "Missing __ROW_ID__ @gnf_data in lifted data Logical"

    # Build inline CSV from the DataFrame.
    # Convert datetime/date columns to ISO 8601 strings before serializing,
    # because pandas to_csv uses space-separated datetime format by default
    # (e.g. "2025-01-01 10:24:10") and the engine expects ISO 8601
    # ("2025-01-01T10:24:10" for DateTime, "2025-01-01" for Date).
    # We use .isoformat() to preserve timezone information.
    df = _prepare_data_for_csv(data_node.data)
    csv_buffer = io.StringIO()
    df.to_csv(csv_buffer, index=False, lineterminator="\n")
    inline_csv = csv_buffer.getvalue()

    # Row ID and CSV index types.
    row_id_lqp_type = meta_type_to_lqp(data_node.vars[0].type)
    csv_index_type = mk_type(lqp_bundle.ir.TypeName.INT)

    # Build column CSV entries and Defs.
    csv_columns, col_defs, first_csv_col_target_id, first_csv_col_value_type = \
        _build_column_defs(ctx, col_gnfs, row_id_lqp_type, csv_index_type)

    # Build the __ROW_ID__ Def.
    row_id_def = _build_row_id_def(
        ctx, row_id_gnf, row_id_lqp_type, csv_index_type,
        first_csv_col_target_id, first_csv_col_value_type,
    )

    # Build the CSVData declaration.
    csv_data = _build_inline_csv_data(inline_csv, csv_columns)

    result: list[lqp_bundle.ir.Declaration] = []
    result.append(csv_data)
    result.append(row_id_def)
    result.extend(col_defs)
    return result


def _prepare_data_for_csv(df: "pd.DataFrame") -> "pd.DataFrame":
    """Convert datetime/date columns to ISO 8601 strings for the engine's CSV parser.

    - Timezone-aware datetimes are converted to UTC (the engine stores DateTime as UTC)
      to avoid non-standard offsets (e.g. +05:03:12) that the CSV parser cannot handle.
    - Uses .isoformat() to produce "T"-separated format (e.g. "2025-01-01T10:24:10+00:00").
    - Date columns (stored as Python date objects) are also converted via .isoformat().
    """
    from datetime import date as _date, datetime as _datetime, timezone as _tz
    result = df.copy()
    for col in result.columns:
        if pd.api.types.is_datetime64_any_dtype(result[col]):
            # Convert tz-aware datetimes to UTC to avoid non-standard tz offsets
            if hasattr(result[col].dt, 'tz') and result[col].dt.tz is not None:
                result[col] = result[col].dt.tz_convert('UTC')
            result[col] = result[col].apply(lambda x: x.isoformat() if pd.notna(x) else "")
        elif result[col].dtype == object:
            sample = result[col].dropna()
            if not sample.empty and isinstance(sample.iloc[0], _datetime):
                result[col] = result[col].apply(
                    lambda x: x.astimezone(_tz.utc).isoformat() if x is not None else ""
                )
            elif not sample.empty and isinstance(sample.iloc[0], _date):
                result[col] = result[col].apply(lambda x: x.isoformat() if x is not None else "")
    return result


def _build_inline_csv_data(
    inline_csv: str,
    gnf_columns: list[lqp_bundle.ir.GNFColumn],
) -> lqp_bundle.ir.CSVData:
    csv_locator = lqp_bundle.ir.CSVLocator(paths=[], inline_data=inline_csv.encode("utf-8"), meta=None)
    csv_config = lqp_bundle.ir.CSVConfig(
        header_row=1,
        skip=0,
        new_line="\n",
        delimiter=",",
        quotechar='"',
        escapechar='"',
        comment="",
        missing_strings=[""],
        decimal_separator=".",
        encoding="utf-8",
        compression="none",
        meta=None,
    )
    return lqp_bundle.ir.CSVData(
        locator=csv_locator,
        config=csv_config,
        columns=gnf_columns,
        asof="",
        meta=None,
    )


def _build_row_id_def(
    ctx: TranslationCtx,
    row_id_gnf: ir.Logical,
    row_id_lqp_type: lqp_bundle.ir.Type,
    csv_index_type: lqp_bundle.ir.Type,
    first_csv_col_target_id: lqp_bundle.ir.RelationId,
    first_csv_col_value_type: lqp_bundle.ir.Type,
) -> lqp_bundle.ir.Def:

    row_id_updates = collect_by_type(ir.Update, row_id_gnf)
    assert len(row_id_updates) == 1
    row_id_update = row_id_updates[0]
    row_id_bindings = _effect_bindings(row_id_update)
    row_id_projection, _ = _translate_bindings(ctx, row_id_bindings)
    row_id_rel_id = get_relation_id(ctx, row_id_update.relation, row_id_projection)

    rid_var = gen_unique_var(ctx, "row_id")
    rid_raw_var = gen_unique_var(ctx, "_row_id")
    rid_dummy_var = gen_unique_var(ctx, "_v")
    rid_shifted_var = gen_unique_var(ctx, "_row_id_shifted")

    two_literal = to_lqp_value(2, types.Int64)
    rid_subtract = mk_primitive("rel_primitive_subtract_monotype", [rid_raw_var, two_literal, rid_shifted_var])
    rid_cast = lqp_bundle.ir.Cast(input=rid_shifted_var, result=rid_var, meta=None)
    rid_atom = lqp_bundle.ir.Atom(name=first_csv_col_target_id, terms=[rid_raw_var, rid_dummy_var], meta=None)

    rid_body = mk_exists(
        [(rid_raw_var, csv_index_type), (rid_dummy_var, first_csv_col_value_type),
         (rid_shifted_var, csv_index_type)],
        mk_and([rid_subtract, rid_cast, rid_atom]),
    )

    return lqp_bundle.ir.Def(
        name=row_id_rel_id,
        body=mk_abstraction([(rid_var, row_id_lqp_type)], rid_body),
        attrs=[],
        meta=None,
    )


def _build_column_defs(
    ctx: TranslationCtx,
    col_gnfs: list[tuple[str, ir.Logical]],
    row_id_lqp_type: lqp_bundle.ir.Type,
    csv_index_type: lqp_bundle.ir.Type,
) -> tuple[list[lqp_bundle.ir.GNFColumn], list[lqp_bundle.ir.Def], lqp_bundle.ir.RelationId, lqp_bundle.ir.Type]:

    csv_columns: list[lqp_bundle.ir.GNFColumn] = []
    col_defs: list[lqp_bundle.ir.Def] = []
    first_csv_col_target_id = None
    first_csv_col_value_type = None

    for col_name, gnf_logical in col_gnfs:
        updates = collect_by_type(ir.Update, gnf_logical)
        assert len(updates) == 1, f"Expected exactly one derive in @gnf_data for column {col_name}"
        update = updates[0]
        derive_rel = update.relation

        bindings = _effect_bindings(update)
        projection, _ = _translate_bindings(ctx, bindings)
        derive_rel_id = get_relation_id(ctx, derive_rel, projection)

        assert len(derive_rel.fields) == 2, f"Expected 2-arity derive relation for column {col_name}"
        col_value_type = meta_type_to_lqp(derive_rel.fields[1].type)

        csv_col_target_id = gen_rel_id(ctx, f"{derive_rel.name}_col_{col_name}")

        if first_csv_col_target_id is None:
            first_csv_col_target_id = csv_col_target_id
            first_csv_col_value_type = col_value_type

        csv_columns.append(lqp_bundle.ir.GNFColumn(
            column_name=col_name,
            target_id=csv_col_target_id,
            types=[csv_index_type, col_value_type],
            meta=None,
        ))

        row_id_var = gen_unique_var(ctx, "row_id")
        col_var = gen_unique_var(ctx, col_name)
        raw_row_id_var = gen_unique_var(ctx, "_row_id")
        shifted_row_id_var = gen_unique_var(ctx, "_row_id_shifted")

        two_literal = to_lqp_value(2, types.Int64)
        subtract_formula = mk_primitive("rel_primitive_subtract_monotype", [raw_row_id_var, two_literal, shifted_row_id_var])
        cast_formula = lqp_bundle.ir.Cast(input=shifted_row_id_var, result=row_id_var, meta=None)
        atom_formula = lqp_bundle.ir.Atom(name=csv_col_target_id, terms=[raw_row_id_var, col_var], meta=None)

        body = mk_exists(
            [(raw_row_id_var, csv_index_type), (shifted_row_id_var, csv_index_type)],
            mk_and([subtract_formula, cast_formula, atom_formula]),
        )

        col_defs.append(lqp_bundle.ir.Def(
            name=derive_rel_id,
            body=mk_abstraction([(row_id_var, row_id_lqp_type), (col_var, col_value_type)], body),
            attrs=[],
            meta=None,
        ))

    assert first_csv_col_target_id is not None, "Expected at least one data column"
    assert first_csv_col_value_type is not None
    return csv_columns, col_defs, first_csv_col_target_id, first_csv_col_value_type


def _translate_to_constraint_decls(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp_bundle.ir.Declaration]:
    constraint_decls: list[lqp_bundle.ir.Declaration] = []
    for task in rule.body:
        if isinstance(task, ir.Logical):
            constraint_decls.extend(_translate_to_constraint_decls(ctx, task))
            continue
        else:
            assert isinstance(task, ir.Require)
            constraint_decls.extend(_translate_to_constraint_decls_from_require(ctx, task))
    return constraint_decls

def _translate_to_constraint_decls_from_require(ctx: TranslationCtx, rule: ir.Require) -> list[lqp_bundle.ir.Declaration]:
    fds = normalized_fds(rule)
    assert len(fds) > 0, f"Expected at least one functional dependency from Require node, got {len(fds)}"

    decls: list[lqp_bundle.ir.Declaration] = []
    for fd in fds:
        # check for unresolved types
        if any(types.is_any(var.type) for var in fd.keys + fd.values):
            warn(f"Ignoring FD with unresolved type: {fd}")
            continue

        lqp_typed_keys = [_translate_term(ctx, key) for key in fd.keys]
        lqp_typed_values = [_translate_term(ctx, value) for value in fd.values]
        lqp_typed_vars:list[Tuple[lqp_bundle.ir.Var, lqp_bundle.ir.Type]] = lqp_typed_keys + lqp_typed_values # type: ignore
        lqp_guard_atoms = [_translate_to_atom(ctx, atom) for atom in fd.guard]
        lqp_guard = mk_abstraction(lqp_typed_vars, mk_and(lqp_guard_atoms))
        lqp_keys:list[lqp_bundle.ir.Var] = [var for (var, _) in lqp_typed_keys] # type: ignore
        lqp_values:list[lqp_bundle.ir.Var] = [var for (var, _) in lqp_typed_values] # type: ignore
        lqp_id = utils.lqp_hash(fd.canonical_str)
        lqp_name:lqp_bundle.ir.RelationId = lqp_bundle.ir.RelationId(id=lqp_id, meta=None)

        decls.append(lqp_bundle.ir.FunctionalDependency(
            name=lqp_name,
            guard=lqp_guard,
            keys=lqp_keys,
            values=lqp_values,
            meta=None
        ))
    return decls

def _translate_algorithms(ctx: TranslationCtx, task: ir.Logical) -> list[lqp_bundle.ir.Declaration]:
    assert is_algorithm_logical(task)
    decls: list[lqp_bundle.ir.Declaration] = []
    for subtask in task.body:
        assert is_algorithm_script(subtask), "Expected all subtasks to be algorithm scripts"
        decls.extend(_translate_algorithm_script(ctx, subtask))
    return decls

def _translate_algorithm_script(ctx: TranslationCtx, alg_task: ir.Sequence) -> list[lqp_bundle.ir.Declaration]:
    assert is_algorithm_script(alg_task), "Expected Sequence @algorithm @script "

    alg_globals = _find_algorithm_global_relation_ids(ctx, alg_task)
    alg_body = _translate_script(ctx, alg_task)

    return [lqp_bundle.ir.Algorithm(global_=alg_globals, body=alg_body, meta=None)]

def _find_algorithm_global_relation_ids(ctx: TranslationCtx, alg_task: ir.Sequence) -> list[lqp_bundle.ir.RelationId]:
    result = []
    updates = collect_by_type(ir.Update, alg_task)
    for update in updates:
        if has_global_annotation(update):
            bindings = _effect_bindings(update)
            projection, _ = _translate_bindings(ctx, bindings)
            rel_id = get_relation_id(ctx, update.relation, projection)
            result.append(rel_id)
    return list(dict.fromkeys(result))

def _translate_script(ctx: TranslationCtx, task: ir.Sequence) -> lqp_bundle.ir.Script:
    assert is_script(task), "Expected a @script Sequence"

    constructs: list[lqp_bundle.ir.Construct] = []

    for subtask in task.tasks:
        if is_empty_instruction(subtask):
            constructs.append(_translate_empty_instruction(ctx, subtask))
        elif isinstance(subtask, ir.Logical):
            constructs.extend(_translate_instruction(ctx, subtask))
        elif isinstance(subtask, ir.Break):
            constructs.append(_translate_break_instruction(ctx, subtask))
        elif is_while_loop(subtask):
            constructs.append(_translate_while_loop(ctx, subtask))
        else:
            raise Exception(f"Unsupported script instruction: {subtask}")

    return lqp_bundle.ir.Script(constructs=constructs, meta=None)

def _translate_while_loop(ctx: TranslationCtx, task: ir.Loop) -> lqp_bundle.ir.Loop:
    assert is_while_loop(task), "Expected a @while Loop"
    assert len(task.iter) == 0, "Temporalized loops not supported"

    while_script_task = task.body
    assert is_while_script(while_script_task), "The body of a @while Loop must be a @while @script Sequence"
    body_script = _translate_script(ctx, while_script_task)

    # No init instructions in the translation of PyRel Loops to to LQP loops
    return lqp_bundle.ir.Loop(init=[], body=body_script, meta=None)

def _translate_break_instruction(ctx: TranslationCtx, task: ir.Break) -> lqp_bundle.ir.Construct:
    body = _translate_to_formula(ctx, task.check)

    ctx.break_rule_counter += 1

    rel_id = gen_rel_id(ctx, "break_cond_" + str(ctx.break_rule_counter))
    return lqp_bundle.ir.Break(
        name = rel_id,
        body = mk_abstraction([], body),
        attrs = [],
        meta = None,
    )

def _translate_empty_instruction(ctx: TranslationCtx, rule: ir.Logical) -> lqp_bundle.ir.Instruction:
    assert is_empty_instruction(rule)
    updates = collect_by_type(ir.Update, rule)
    assert len(updates) == 1
    update = updates[0]
    bindings = _effect_bindings(update)

    # We need  to make sure that variable names have a leading underscore
    normalized_bindings:Sequence[ir.Var] = []
    for v in bindings:
        assert isinstance(v, ir.Var)
        if not v.name.startswith("_"):
            v = ir.Var(v.type, "_" + v.name)
        normalized_bindings.append(v)

    projection, eqs = _translate_bindings(ctx, normalized_bindings)
    assert len(eqs) == 0
    rel_id = get_relation_id(ctx, update.relation, projection)
    abstraction = mk_abstraction(projection, mk_or([])) # empty body = false
    return lqp_bundle.ir.Assign(name = rel_id, body = abstraction, attrs = [], meta = None)

def _translate_instruction(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp_bundle.ir.Instruction]:
    effects = collect_by_type((ir.Update, ir.Output), rule)

    if len(effects) == 0:
        return []

    conjuncts = _translate_to_formula(ctx, rule)
    res = []
    for effect in effects:
        assert isinstance(effect, ir.Update), f"Got an effect of type {type(effect)} in a loop, which is invalid."

        bindings = _effect_bindings(effect)
        projection, eqs = _translate_bindings(ctx, bindings)

        eqs.append(conjuncts)
        new_body = mk_and(eqs)

        rel_id = get_relation_id(ctx, effect.relation, projection)
        abstraction = mk_abstraction(projection, new_body)

        annotations = effect.annotations | effect.relation.annotations
        attrs = _translate_annotations(annotations)

        upsert = get_upsert_annotation(effect)
        monoid = get_monoid_annotation(effect)
        monus = get_monus_annotation(effect)

        if has_assign_annotation(effect):
            res.append(lqp_bundle.ir.Assign(name = rel_id, body = abstraction, attrs = attrs, meta = None))
        elif upsert is not None:
            res.append(lqp_bundle.ir.Upsert(value_arity=get_arity(upsert), name = rel_id, body = abstraction, attrs = attrs, meta = None))
        elif monoid is not None:
            res.append(lqp_bundle.ir.MonoidDef(
                value_arity=get_arity(monoid),
                monoid=construct_monoid(monoid),
                name = rel_id,
                body = abstraction,
                attrs = attrs,
                meta = None
            ))
        elif monus is not None:
            res.append(lqp_bundle.ir.MonusDef(
                value_arity=get_arity(monus),
                monoid=construct_monoid(monus),
                name = rel_id,
                body = abstraction,
                attrs = attrs,
                meta = None
            ))

    return res

def _translate_to_standard_decl(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp_bundle.ir.Declaration]:
    effects = collect_by_type((ir.Output, ir.Update), rule)

    # TODO: should this ever actually come in as input?
    if len(effects) == 0:
        return []

    conjuncts = _translate_to_formula(ctx, rule)
    return [_translate_effect(ctx, effect, conjuncts) for effect in effects]

def _translate_annotations(annotations: FrozenOrderedSet[ir.Annotation]) -> list[lqp_bundle.ir.Attribute]:
    attributes = []
    for annotation in annotations:
        if annotation.relation.name in supported_lqp_annotations:
            attr = _translate_annotation(annotation)
            if attr is not None:
                attributes.append(attr)
    return attributes

def _translate_annotation(annotation: ir.Annotation) -> Optional[lqp_bundle.ir.Attribute]:
    if any(not isinstance(a, ir.Literal) for a in annotation.args):
        warn("LQP currently ignores annotation parameters with non-literal values")
        return None
    if annotation.relation.name == builtins.incremental_maintenance.name:
        return _translate_ivm_annotation(annotation)
    else:
        # Convert literal arguments to LQP values
        args = []
        for a in annotation.args:
            assert isinstance(a, ir.Literal)
            args.append(mk_value(a.value))
        return mk_attribute(annotation.relation.name, args)

def _translate_ivm_annotation(annotation: ir.Annotation) -> Optional[lqp_bundle.ir.Attribute]:
    if len(annotation.args) == 0:
        warn("incremental_maintenance annotation requires a value argument.")
        return None
    arg = annotation.args[0]
    assert isinstance(arg, ir.Literal)
    value = arg.value
    if value not in ("all", "off"):
        warn(f"Invalid value for incremental_maintenance annotation: {value}. Expected 'all' or 'off'.")
        return None
    else:
        if value == "all":
            return mk_attribute("must_try_maintain", [])
        elif value == "off":
            return mk_attribute("nomaintain", [])

# Translates an effect (export, output, or update) into the corresponding def. Note that
# this method only generates the def and not the LQP read operation, which is added later by
# `_get_output_reads` and `_get_export_reads`.
def _translate_effect(ctx: TranslationCtx, effect: Union[ir.Output, ir.Update], body: lqp_bundle.ir.Formula) -> lqp_bundle.ir.Declaration:
    bindings = _effect_bindings(effect)

    if isinstance(effect, ir.Output):
        projection, eqs, suffix = _translate_output_bindings(ctx, bindings)
        meta_id = effect.id

        if helpers.is_export(effect):
            def_name = "export_relation" + suffix
        else:
            def_name = "output" + suffix

        def_name = ctx.output_names.get_name_by_id(meta_id, def_name)
        rel_id = get_output_id(ctx, def_name, meta_id)
    else:
        projection, eqs = _translate_bindings(ctx, bindings)
        def_name = effect.relation.name
        rel_id = get_relation_id(ctx, effect.relation, projection)

    eqs.append(body)
    new_body = mk_and(eqs)

    # Context bookkeeping for exports and outputs
    if helpers.is_export(effect):
        # The row id is the first n-1 elements, and the actual data is the last element. Its
        # type is stored in the first element of the tuple.
        col_type = projection[-1][1]
        _col_num_match = re.search(r"export_relation_col([0-9]+)", def_name)
        assert _col_num_match, f"Could not find column number in suffix: {def_name}"
        col_num = int(_col_num_match.group(1))
        col_name = f"col{col_num}"
        if isinstance(effect, ir.Output) and len(effect.aliases) > 0:
            aliases_list = list(effect.aliases)
            col_name = aliases_list[-1][0]

        # Extract export FQN from annotations
        export_fqn_str = extract_export_fqn(effect.annotations)

        ctx.export_descriptors.append(ExportDescriptor(
            relation_id=rel_id,
            column_name=col_name,
            column_number=col_num,
            column_type=col_type,
            export_fqn=export_fqn_str,
        ))
    elif isinstance(effect, ir.Output):
        ctx.output_ids.append((rel_id, def_name))

    # First we collect annotations on the effect itself, e.g. from something like
    # `select(...).annotate(...)`.
    annotations = effect.annotations
    if isinstance(effect, ir.Update):
        # Then we translate annotations on the relation itself, e.g.
        # ```
        # Bar.foo = model.Relationship(...)
        # Bar.foo.annotate(...)
        # ```
        annotations = annotations | effect.relation.annotations

    return lqp_bundle.ir.Def(
        name = rel_id,
        body = mk_abstraction(projection, new_body),
        attrs = _translate_annotations(annotations),
        meta = None,
    )

def _translate_output_bindings(ctx: TranslationCtx, bindings: list[ir.Value]) -> Tuple[list[Tuple[lqp_bundle.ir.Var, lqp_bundle.ir.Type]], list[lqp_bundle.ir.Formula], str]:
    symbol_literals = []
    non_symbols = []
    for binding in bindings:
        if isinstance(binding, ir.Literal) and binding.type == types.Symbol:
            symbol_literals.append(binding.value)
        else:
            non_symbols.append(binding)
    projection, eqs = _translate_bindings(ctx, non_symbols)
    if len(symbol_literals) > 0:
        name_suffix = "_"
        name_suffix += "_".join(symbol_literals)
    else:
        name_suffix = ""

    return projection, eqs, name_suffix

def _translate_rank(ctx: TranslationCtx, rank: ir.Rank, body: lqp_bundle.ir.Formula) -> lqp_bundle.ir.Formula:
    # Ascending rank is constructed using rel_primitive_sort. If a limit is added to an
    # ascending rank we can use rel_primitive_top for an efficient evaluation.
    #
    # Descending rank is constructed as an ascending sort, then the ascending rank is
    # subtracted from a count of the elements (plus 1 so that we still start at 1).
    # Adding a limit to a descending rank is done by adding a filter of rank <= limit.

    # Limits are the sort plus a filter on rank <= limit.
    if all(o for o in rank.arg_is_ascending):
        ascending = True
    elif all(not o for o in rank.arg_is_ascending):
        ascending = False
    else:
        raise Exception("Mixed orderings in rank are not supported yet.")

    # Filter out the group-by variables, since they are introduced outside the rank.
    input_args, input_eqs = _translate_bindings(ctx, list(rank.args))
    introduced_meta_projs = [arg for arg in rank.projection if arg not in rank.group and arg not in rank.args]
    projected_args, projected_eqs = _translate_bindings(ctx, list(introduced_meta_projs))

    # rank expects an Int128 result, but the primitive will return an Int64 result.
    # we need to set up an intermediary variable to hold the Int64 result, and a cast
    # to convert it to Int128.
    result_var, _ = _translate_term(ctx, rank.result)
    # The primitive will return an Int64 result, so we need a var to hold the intermediary.
    result_64_var = gen_unique_var(ctx, "v_rank")
    result_64_type = mk_type(lqp_bundle.ir.TypeName.INT)

    cast = lqp_bundle.ir.Cast(input=result_64_var, result=result_var, meta=None)

    body = mk_and([body] + input_eqs + projected_eqs)
    abstr_args = input_args + projected_args

    if ascending:
        ranker = _translate_ascending_rank(ctx, rank.limit, result_64_var, body, abstr_args)
    else:
        ranker = _translate_descending_rank(ctx, rank.limit, result_64_var, body, abstr_args)

    return mk_exists([(result_64_var, result_64_type)], mk_and([ranker, cast]))

def _translate_descending_rank(ctx: TranslationCtx, limit: int, result: lqp_bundle.ir.Var, body: lqp_bundle.ir.Formula, abstr_args) -> lqp_bundle.ir.Formula:
    result_var = result
    result_type = mk_type(lqp_bundle.ir.TypeName.INT)

    # Rename abstracted args in the body to new variable names
    var_map = {var.name: gen_unique_var(ctx, 't_' + var.name) for (var, _) in abstr_args}
    body = utils.rename_vars_formula(body, var_map)
    new_abstr_args = [(var_map[var.name], typ) for (var, typ) in abstr_args]

    # Construct a conjunction of the ranking, a counter for the body, a subtraction
    # of the rank from the count and an addition of 1. Wrap this in an abstraction.
    count_res = gen_unique_var(ctx, "count_res")

    # Add one to the count to account for the rank starting at 1.
    one, one_eq = constant_to_var(ctx, to_lqp_value(1, types.Int64), "one")
    one_bigger = gen_unique_var(ctx, "one_bigger")
    addition = mk_primitive("rel_primitive_add_monotype", [count_res, one, one_bigger])

    # Subtract the rank from the count + 1
    asc_rank = gen_unique_var(ctx, "asc_rank")
    subtraction = mk_primitive("rel_primitive_subtract_monotype", [one_bigger, result_var, asc_rank])

    # Construct the ranking
    desc_ranking_terms = [asc_rank] + [v[0] for v in abstr_args]
    ranking = lqp_bundle.ir.FFI(
        meta=None,
        name="rel_primitive_sort",
        args=[mk_abstraction(new_abstr_args, body)],
        terms=desc_ranking_terms,
    )

    # Count the number of rows in the body
    count_type = meta_type_to_lqp(types.Int64)
    count_var, count_eq = constant_to_var(ctx, to_lqp_value(1, types.Int64), "counter")
    desc_body = mk_and([body, count_eq])
    aggr_abstr_args = new_abstr_args + [(count_var, count_type)]
    count_aggr = lqp_bundle.ir.Reduce(
        op=lqp_operator(
            ctx,
            "count",
            "count",
            mk_type(lqp_bundle.ir.TypeName.INT)
        ),
        body=mk_abstraction(aggr_abstr_args, desc_body),
        terms=[count_res],
        meta=None
    )

    # Bring it all together and do the maths.
    ranking = mk_exists(
        vars=[
            (asc_rank, result_type),
            (count_res, result_type),
            (one, result_type),
            (one_bigger, result_type)
        ],
        value=mk_and([ranking, count_aggr, one_eq, addition, subtraction])
    )

    # If there is a limit, we need to add a filter to the ranking.
    # Wrap with a rank <= limit
    if limit != 0:
        limit_term, _ = _translate_term(ctx, ir.Literal(types.Int64, limit))
        limiter = mk_primitive("rel_primitive_lt_eq_monotype", [result_var, limit_term])
        ranking = mk_and([ranking, limiter])

    return ranking

def _translate_ascending_rank(ctx: TranslationCtx, limit: int, result_var: lqp_bundle.ir.Var, body: lqp_bundle.ir.Formula, abstr_args) -> lqp_bundle.ir.Formula:
    terms = [result_var] + [v[0] for v in abstr_args]

    # Rename abstracted args in the body to new variable names
    var_map = {var.name: gen_unique_var(ctx, 't_' + var.name) for (var, _) in abstr_args}
    body = utils.rename_vars_formula(body, var_map)
    new_abstr_args = [(var_map[var.name], typ) for (var, typ) in abstr_args]
    sort_abstr = mk_abstraction(new_abstr_args, body)

    if limit == 0:
        return lqp_bundle.ir.FFI(
            meta=None,
            name="rel_primitive_sort",
            args=[sort_abstr],
            terms=terms,
        )
    else:
        limit_type = meta_type_to_lqp(types.Int64)
        limit_var, limit_eq = constant_to_var(ctx, to_lqp_value(limit, types.Int64), "limit")
        limit_abstr = mk_abstraction([(limit_var, limit_type)], limit_eq)
        return lqp_bundle.ir.FFI(
            meta=None,
            name="rel_primitive_top",
            args=[sort_abstr, limit_abstr],
            terms=terms,
        )

def _rename_shadowed_abstraction_vars(
    ctx: TranslationCtx,
    aggr: ir.Aggregate,
    abstr_args: list[Tuple[lqp_bundle.ir.Var, lqp_bundle.ir.Type]],
    body_conjs: list[lqp_bundle.ir.Formula]
) -> list[Tuple[lqp_bundle.ir.Var, lqp_bundle.ir.Type]]:
    """
    Rename abstraction variables that shadow group-by variables.

    This can happen when the same variable appears in both aggr.group and as an input
    to the aggregation, e.g., min(Person.age).per(Person.age). The group-by variables
    are in the outer scope, while the abstraction parameters are in the inner scope,
    so we need different names to avoid shadowing.
    """
    # Get the LQP names of group-by variables
    group_var_names = set()
    for group_var in aggr.group:
        lqp_var = _translate_var(ctx, group_var)
        group_var_names.add(lqp_var.name)

    # Rename any abstraction parameters that conflict with group-by variables
    renamed_abstr_args = []
    for var, typ in abstr_args:
        if var.name in group_var_names:
            # This variable shadows a group-by variable, so rename it
            fresh_var = gen_unique_var(ctx, var.name)
            # Add an equality constraint: fresh_var == var
            # var is a free variable referring to the outer scope group-by variable
            body_conjs.append(mk_primitive("rel_primitive_eq", [fresh_var, var]))
            renamed_abstr_args.append((fresh_var, typ))
        else:
            renamed_abstr_args.append((var, typ))

    return renamed_abstr_args

def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp_bundle.ir.Formula) -> Union[lqp_bundle.ir.Reduce, lqp_bundle.ir.Formula]:
    # TODO: handle this properly
    aggr_name = aggr.aggregation.name
    supported_aggrs = ("sum", "count", "avg", "min", "max", "rel_primitive_solverlib_ho_appl")
    assert aggr_name in supported_aggrs, f"only support {supported_aggrs} for now, not {aggr.aggregation.name}"

    meta_output_terms = []
    meta_input_terms = []

    for (field, arg) in zip(aggr.aggregation.fields, aggr.args):
        if field.input:
            meta_input_terms.append(arg)
        else:
            meta_output_terms.append(arg)

    output_terms = [_translate_term(ctx, term) for term in meta_output_terms]
    output_vars = [term[0] for term in output_terms]

    body_conjs = [body]
    input_args, input_eqs = _translate_bindings(ctx, meta_input_terms)

    # TODO: Can this safely be applied to all aggregates?
    if aggr_name in ("sum", "min", "max"):
        assert len(output_terms) == 1, f"{aggr_name} expects a single output variable"
        assert len(meta_input_terms) == 1, f"{aggr_name} expects a single input variable"
        assert isinstance(meta_output_terms[0], ir.Var)
        assert input_args[0][1] == output_terms[0][1], f"{aggr_name}({input_args[0][1].type_name}) had output type of {output_terms[0][1].type_name}"

    # Filter out the group-by variables, since they are introduced outside the aggregation.
    # Input terms are added later below.
    introduced_meta_projs = [arg for arg in aggr.projection if arg not in aggr.group and arg not in meta_input_terms]
    projected_args, projected_eqs = _translate_bindings(ctx, list(introduced_meta_projs))
    body_conjs.extend(input_eqs)
    body_conjs.extend(projected_eqs)
    abstr_args: list[Tuple[lqp_bundle.ir.Var, lqp_bundle.ir.Type]] = projected_args + input_args

    # Rename abstraction variables that shadow group-by variables
    abstr_args = _rename_shadowed_abstraction_vars(ctx, aggr, abstr_args, body_conjs)

    if aggr_name == "count":
        assert len(output_terms) == 1, "Count and avg expect a single output variable"
        assert isinstance(meta_output_terms[0], ir.Var)
        # Count sums up "1" for each row. We use the expected output type for the type
        # of the count variable.
        typ = meta_type_to_lqp(meta_output_terms[0].type)
        one_var, eq = constant_to_var(ctx, to_lqp_value(1, meta_output_terms[0].type), "one")
        body_conjs.append(eq)
        abstr_args.append((one_var, typ))

    # Average needs to wrap the reduce in Exists(Conjunction(Reduce, div))
    if aggr_name == "avg":
        assert len(output_vars) == 1, "avg should only have one output variable"
        output_var = output_vars[0]

        # Count sums up "1" for each row. We make the reasonably safe assumption that there
        # are less than 2^31 rows in the body.
        count_type = meta_type_to_lqp(types.Int64)
        one_var, eq = constant_to_var(ctx, to_lqp_value(1, types.Int64), "one")
        body_conjs.append(eq)
        abstr_args.append((one_var, count_type))

        # The average will produce two output variables: sum and count.
        sum_result = gen_unique_var(ctx, "sum")
        count_result = gen_unique_var(ctx, "count")

        # Second to last is the variable we're summing over.
        (sum_var, sum_type) = abstr_args[-2]
        body = mk_and(body_conjs)

        result = lqp_bundle.ir.Reduce(
            op=lqp_avg_op(ctx, aggr.aggregation.name, sum_var.name, sum_type),
            body=mk_abstraction(abstr_args, body),
            terms=[sum_result, count_result],
            meta=None,
        )

        if sum_type == count_type:
            div = mk_primitive("rel_primitive_divide_monotype", [sum_result, count_result, output_var])
            conjunction = mk_and([result, div])

            # Finally, we need to wrap everything in an `exists` to project away the sum and
            # count variables and only keep the result of the division.
            result_terms = [(sum_result, sum_type), (count_result, count_type)]
        else:
            # If the sum type and count type don't match, we need to cast the count
            count_casted = gen_unique_var(ctx, "count_casted")
            count_cast = lqp_bundle.ir.Cast(input=count_result, result=count_casted, meta=None)

            div = mk_primitive("rel_primitive_divide_monotype", [sum_result, count_casted, output_var])
            conjunction = mk_and([result, count_cast, div])

            # Finally, we need to wrap everything in an `exists` to project away the sum and
            # count variables and only keep the result of the division.
            result_terms = [(sum_result, sum_type), (count_result, count_type), (count_casted, sum_type)]

        return mk_exists(result_terms, conjunction)

    # `input_args` hold the types of the input arguments, but they may have been modified
    # if we're dealing with a count, so we use `abstr_args` to find the type.
    (aggr_arg, aggr_arg_type) = abstr_args[-1]
    body = mk_and(body_conjs)

    # Group-bys do not need to be handled at all, since they are introduced outside already
    reduce = lqp_bundle.ir.Reduce(
        op=lqp_operator(ctx, aggr.aggregation.name, aggr_arg.name, aggr_arg_type),
        body=mk_abstraction(abstr_args, body),
        terms=output_vars,
        meta=None
    )
    return reduce

def _translate_to_formula(ctx: TranslationCtx, task: ir.Task) -> lqp_bundle.ir.Formula:
    if isinstance(task, ir.Logical):
        # For aggregates and ranks, the expected format is:
        #
        # Logical
        #   body_task1
        #   body_task2
        #   ...
        #   aggregate/rank task
        #
        # If we see that the Logical is in this format, it should be translated as an
        # aggregate/rank node.
        groups = group_tasks(task.body, {
            "aggregates": ir.Aggregate,
            "ranks": ir.Rank,
        })

        aggregates = groups.get("aggregates", OrderedSet[ir.Task]())
        ranks = groups.get("ranks", OrderedSet[ir.Task]())

        if aggregates or ranks:
            conjuncts = []
            body = mk_and([_translate_to_formula(ctx, t) for t in task.body])
            for aggr in aggregates:
                assert(isinstance(aggr, ir.Aggregate))
                conjuncts.append(_translate_aggregate(ctx, aggr, body))
            for rank in ranks:
                assert(isinstance(rank, ir.Rank))
                conjuncts.append(_translate_rank(ctx, rank, body))
        else:
            # If there are no aggregates or ranks, translate as a normal conjunction
            conjuncts = [_translate_to_formula(ctx, child) for child in task.body]
        return mk_and(conjuncts)
    elif isinstance(task, ir.Lookup):
        return _translate_to_atom(ctx, task)
    elif isinstance(task, ir.Not):
        return lqp_bundle.ir.Not(arg=_translate_to_formula(ctx, task.task), meta=None)
    elif isinstance(task, ir.Exists):
        lqp_vars, conjuncts = _translate_bindings(ctx, list(task.vars))
        conjuncts.append(_translate_to_formula(ctx, task.task))
        return mk_exists(lqp_vars, mk_and(conjuncts))
    elif isinstance(task, ir.Construct):
        assert len(task.values) >= 1, "Construct should have at least one value"
        terms = [_translate_term(ctx, arg) for arg in task.values]
        result_term = _translate_term(ctx, task.id_var)
        terms.append(result_term)
        assert result_term[1].type_name == lqp_bundle.ir.TypeName.UINT128, \
            f"Attempting to store a {task.id_var} in a type `{result_term[1].type_name}`"

        return mk_primitive("rel_primitive_hash_tuple_uint128", [v for v, _ in terms])
    elif isinstance(task, ir.Union):
        disjs = [_translate_to_formula(ctx, child) for child in task.tasks]
        return mk_or(disjs)
    elif isinstance(task, (ir.Output, ir.Update)):
        # Nothing to do here, handled in _translate_to_decls
        return mk_and([])
    elif isinstance(task, (ir.Aggregate, ir.Rank)):
        # Nothing to do here, handled at the Logical level
        return mk_and([])
    elif isinstance(task, ir.Data):
        # Data nodes should be lifted before translation to LQP
        raise NotImplementedError("Data nodes should be lifted before translation to LQP")
    else:
        raise NotImplementedError(f"Unknown task type (formula): {type(task)}")

# Only used for translating terms on relatoms, which can be specialized values.
def _translate_relterm(ctx: TranslationCtx, term: ir.Value) -> Tuple[lqp_bundle.ir.RelTerm, lqp_bundle.ir.Type]:
    if isinstance(term, ir.Literal) and term.type == types.Symbol:
        if isinstance(term.value, str):
            value = mk_value(term.value)
            return mk_specialized_value(value), meta_type_to_lqp(types.String)
        elif isinstance(term.value, int):
            value = mk_value(term.value)
            return mk_specialized_value(value), meta_type_to_lqp(types.Int64)
        else:
            raise NotImplementedError(f"Cannot specialize literal of type {type(term.value)}")
    return _translate_term(ctx, term)

def _translate_term(ctx: TranslationCtx, term: ir.Value) -> Tuple[lqp_bundle.ir.Term, lqp_bundle.ir.Type]:
    if isinstance(term, ir.ScalarType):
        # TODO: ScalarType is not like other terms, should be handled separately.
        return to_lqp_value(term.name, types.String), meta_type_to_lqp(types.String)
    elif isinstance(term, ir.Var):
        t = meta_type_to_lqp(term.type)
        return _translate_var(ctx, term), t
    else:
        assert isinstance(term, ir.Literal), f"Cannot translate value {term!r} of type {type(term)} to LQP Term; neither Var nor Literal."
        v = to_lqp_value(term.value, term.type)
        return v, meta_type_to_lqp(term.type)

def _translate_to_atom(ctx: TranslationCtx, task: ir.Lookup) -> lqp_bundle.ir.Formula:
    if task.relation == builtins.cast:
        assert len(task.args) == 3, f"expected three terms for {task.relation.name}, got {len(task.args)}"

        terms = []
        for arg in task.args[1:]:
            term, _ = _translate_relterm(ctx, arg)
            terms.append(term)
        return lqp_bundle.ir.Cast(input=terms[0], result=terms[1], meta=None)

    if task.relation == builtins.join:
        return _translate_join(ctx, task)
    elif task.relation == builtins.infomap:
        return _translate_infomap(ctx, task)
    elif task.relation == builtins.louvain:
        return _translate_louvain(ctx, task)
    elif task.relation == builtins.label_propagation:
        return _translate_label_propagation(ctx, task)

    terms = []
    term_types = []
    for arg in task.args:
        # Handle varargs, which come wrapped in a tuple.
        if isinstance(arg, tuple):
            for vararg in arg:
                term, ty = _translate_relterm(ctx, vararg)
                terms.append(term)
                term_types.append(ty)
        else:
            term, ty = _translate_relterm(ctx, arg)
            terms.append(term)
            term_types.append(ty)

    if builtins.is_pragma(task.relation):
        lqp_name = pragma_to_lqp_name(task.relation.name)
        return mk_pragma(lqp_name, terms)

    if builtins.is_builtin(task.relation):
        return build_primitive(task.relation.name, terms, term_types)

    if helpers.is_external(task.relation):
        return lqp_bundle.ir.RelAtom(name=task.relation.name, terms=terms, meta=None)

    projection, _ = _translate_bindings(ctx, list(task.args))
    rid = get_relation_id(ctx, task.relation, projection)
    return lqp_bundle.ir.Atom(name=rid, terms=terms, meta=None)


def get_relation_id(ctx: TranslationCtx, relation: ir.Relation, projection: list[Tuple[lqp_bundle.ir.Var, lqp_bundle.ir.Type]] = []) -> lqp_bundle.ir.RelationId:
    types = "_".join([str(t.type_name) for (_, t) in projection])
    if types:
        types = "_" + types
    if relation.id in ctx.def_names.id_to_name:
        unique_name = ctx.def_names.id_to_name[relation.id]
    else:
        name = helpers.python_qualified_name(relation)
        unique_name = ctx.def_names.get_name_by_id(relation.id, name)

    return utils.gen_rel_id(ctx, unique_name, types)

def get_output_id(ctx: TranslationCtx, orig_name: str, metamodel_id: int) -> lqp_bundle.ir.RelationId:
    unique_name = ctx.output_names.get_name_by_id(metamodel_id, orig_name)
    return utils.gen_rel_id(ctx, unique_name)

def _translate_bindings(ctx: TranslationCtx, bindings: Sequence[ir.Value]) -> Tuple[list[Tuple[lqp_bundle.ir.Var, lqp_bundle.ir.Type]], list[lqp_bundle.ir.Formula]]:
    lqp_vars = []
    conjuncts = []
    for binding in bindings:
        lqp_var, typ, eq = binding_to_lqp_var(ctx, binding)
        lqp_vars.append((lqp_var, typ))
        if eq is not None:
            conjuncts.append(eq)

    return lqp_vars, conjuncts

def binding_to_lqp_var(ctx: TranslationCtx, binding: ir.Value) -> Tuple[lqp_bundle.ir.Var, lqp_bundle.ir.Type, Union[None, lqp_bundle.ir.Formula]]:
    if isinstance(binding, ir.Var):
        var, typ = _translate_term(ctx, binding)
        assert isinstance(var, lqp_bundle.ir.Var)
        return var, typ, None
    elif isinstance(binding, ir.Literal):
        lqp_value = to_lqp_value(binding.value, binding.type)
        var, formula = constant_to_var(ctx, lqp_value)
        return var, meta_type_to_lqp(binding.type), formula
    else:
        raise Exception(f"Unsupported binding type: {type(binding)}")

def to_lqp_value(value: ir.PyValue, value_type: ir.Type) -> lqp_bundle.ir.Value:
    typ = meta_type_to_lqp(value_type)

    # Ensure int values match the requested integer type.
    if typ.type_name == lqp_bundle.ir.TypeName.INT and isinstance(value, int):
        assert lqp_types.INT_MIN <= value <= lqp_types.INT_MAX, f"{value} out of range for a 64-bit INT value"
        val = value
    elif typ.type_name == lqp_bundle.ir.TypeName.INT128 and isinstance(value, int):
        assert lqp_types.INT128_MIN <= value <= lqp_types.INT128_MAX, f"{value} out of range for an INT128 value"
        val = lqp_bundle.ir.Int128Value(value=value, meta=None)
    elif typ.type_name == lqp_bundle.ir.TypeName.UINT128 and isinstance(value, int):
        assert lqp_types.UINT128_MIN <= value <= lqp_types.UINT128_MAX, f"{value} out of range for an UINT128 value"
        val = lqp_bundle.ir.UInt128Value(value=value, meta=None)
    elif typ.type_name == lqp_bundle.ir.TypeName.FLOAT and isinstance(value, float):
        val = value
    elif typ.type_name == lqp_bundle.ir.TypeName.STRING and isinstance(value, str):
        val = value
    elif typ.type_name == lqp_bundle.ir.TypeName.DECIMAL and isinstance(value, (int, float, PyDecimal)):
        precision = typ.parameters[0].value
        scale = typ.parameters[1].value
        assert isinstance(precision, int) and isinstance(scale, int)
        val = lqp_bundle.ir.DecimalValue(precision=precision, scale=scale, value=PyDecimal(value), meta=None)
    elif typ.type_name == lqp_bundle.ir.TypeName.DATE and isinstance(value, date):
        val = lqp_bundle.ir.DateValue(value=value, meta=None)
    elif typ.type_name == lqp_bundle.ir.TypeName.DATETIME and isinstance(value, datetime):
        utc_value = value.astimezone(timezone.utc) if value.tzinfo is not None else value # Convert to UTC cf. Iceberg
        val = lqp_bundle.ir.DateTimeValue(value=utc_value, meta=None)
    elif typ.type_name == lqp_bundle.ir.TypeName.BOOLEAN and isinstance(value, bool):
        val = lqp_bundle.ir.BooleanValue(value=value, meta=None)
    else:
        raise Exception(f"Unsupported type for constant: {typ.type_name} with value {value} of type {type(value)}")

    return mk_value(val)

def constant_to_var(ctx: TranslationCtx, value: lqp_bundle.ir.Value, name_hint: str = "cvar") -> Tuple[lqp_bundle.ir.Var, lqp_bundle.ir.Formula]:
    var = gen_unique_var(ctx, name_hint)
    eq = mk_primitive("rel_primitive_eq", [var, value])
    return var, eq

def _extract_pyrel_error_ids(ctx: TranslationCtx, model: ir.Model) -> list[Tuple[lqp_bundle.ir.RelationId, str]]:
    effects = collect_by_type(ir.Update, model)
    pyrel_error_attrs = []
    # We use a separate counter here to avoid modifying the current model.
    i = 1
    for effect in effects:
        if "pyrel_error_attrs" not in effect.relation.name:
            continue
        projection, _ = _translate_bindings(ctx, _effect_bindings(effect))
        rel_id = get_relation_id(ctx, effect.relation, projection)
        name = f"{effect.relation.name}_{i}"
        pyrel_error_attrs.append((rel_id, name))
        i += 1

    return pyrel_error_attrs

# Translate a relation reference into an abstraction over its fields.
def _translate_relation_ref(ctx: TranslationCtx, relation: ir.Relation) -> lqp_bundle.ir.Abstraction:
    projection = []
    for field in relation.fields:
        var = gen_unique_var(ctx, field.name)
        typ = meta_type_to_lqp(field.type)
        projection.append((var, typ))

    rid = get_relation_id(ctx, relation, projection)
    atom = lqp_bundle.ir.Atom(name=rid, terms=[var for (var, _) in projection], meta=None)
    return mk_abstraction(projection, atom)

# Common translation logic for graph algorithms.
# task.args[0] : normalized weight list (int64, int64, float)
# task.args[1] : normalized node count (relation or constant int64)
# task.args[2] : normalized edge count (relation or constant int64)
# task.args[3:-3] : algorithm parameters (var or constant)
# task.args[-3] : diagnostic info
# task.args[-2] : node index
# task.args[-1] : community ident
def _translate_graph_common(name: str, ctx: TranslationCtx, task: ir.Lookup):
    abstractions = []

    assert isinstance(task.args[0], ir.Relation), \
        f"Expected relation as first arg to {name}, got {task.args[0]}:{type(task.args[0])}"
    abstractions.append(_translate_relation_ref(ctx, task.args[0]))

    # Allow constant args for node and edge count
    for arg in task.args[1:3]:
        if isinstance(arg, ir.Relation):
            abst = _translate_relation_ref(ctx, arg)
            typ = abst.vars[0][1]
            assert typ.type_name == lqp_bundle.ir.TypeName.INT, \
                f"Expected Int64 types for node and edge counts, got type {typ.type_name}"
            abstractions.append(abst)
        else:
            var, typ, eq = binding_to_lqp_var(ctx, arg)
            assert eq is not None, \
                f"Expected equality formula for {name} arg {arg}:{type(arg)}"
            abstractions.append(mk_abstraction([(var, typ)], mk_and([eq])))

    for arg in task.args[3:-3]:
        var, typ, eq = binding_to_lqp_var(ctx, arg)
        if eq:
            abstractions.append(mk_abstraction([(var, typ)], mk_and([eq])))
        else:
            print(f"Primitive graph algorithm arg without eq:\n var:{var}, typ:{typ}\n arg:{arg}:{type(arg)}")
            abstractions.append(mk_abstraction([(var, typ)], mk_and([])))

    terms = []
    for arg in task.args[-3:]:
        term, _ = _translate_relterm(ctx, arg)
        terms.append(term)

    return (abstractions, terms)

def _translate_infomap(ctx: TranslationCtx, task: ir.Lookup) -> lqp_bundle.ir.Formula:
    abstractions, terms = _translate_graph_common("infomap", ctx, task)

    return lqp_bundle.ir.FFI(
        meta=None,
        name="rel_primitive_infomap",
        args=abstractions,
        terms=terms,
    )

def _translate_louvain(ctx: TranslationCtx, task: ir.Lookup) -> lqp_bundle.ir.Formula:
    abstractions, terms = _translate_graph_common("louvain", ctx, task)

    return lqp_bundle.ir.FFI(
        meta=None,
        name="rel_primitive_louvain",
        args=abstractions,
        terms=terms,
    )

def _translate_label_propagation(ctx: TranslationCtx, task: ir.Lookup) -> lqp_bundle.ir.Formula:
    abstractions, terms = _translate_graph_common("label_propagation", ctx, task)

    return lqp_bundle.ir.FFI(
        meta=None,
        name="rel_primitive_async_label_propagation",
        args=abstractions,
        terms=terms,
    )

# Hard-coded implementation of Rel's string_join
def _translate_join(ctx: TranslationCtx, task: ir.Lookup) -> lqp_bundle.ir.Formula:
    assert len(task.args) == 3
    (strs, separator, target) = task.args
    assert isinstance(separator, ir.Node)
    assert isinstance(target, ir.Var)
    assert isinstance(strs, tuple)

    # Enumerate the strings we're going to be reducing over
    enumerated_conjunctions = []
    index_var, index_type = gen_unique_var(ctx, "idx"), meta_type_to_lqp(types.Int64)
    string_var, string_type = gen_unique_var(ctx, "s"), meta_type_to_lqp(types.String)
    for i, s in enumerate(strs):
        index = to_lqp_value(i, types.Int64)
        string = _translate_term(ctx, s)[0]
        eq_idx = mk_primitive("rel_primitive_eq", [index_var, index])
        eq_string = mk_primitive("rel_primitive_eq", [string_var, string])
        enumerated_conjunctions.append(mk_and([eq_idx, eq_string]))
    body = mk_abstraction([(index_var, index_type), (string_var, string_type)], mk_or(enumerated_conjunctions))

    # Make the function we're reducing
    a_var, a_type = gen_unique_var(ctx, "a"), meta_type_to_lqp(types.String)
    b_var, b_type = gen_unique_var(ctx, "b"), meta_type_to_lqp(types.String)
    curr_var, curr_type = gen_unique_var(ctx, "curr"), meta_type_to_lqp(types.String)
    res_var, res_type = gen_unique_var(ctx, "res"), meta_type_to_lqp(types.String)

    sep, _ = _translate_term(ctx, separator)

    op = mk_abstraction(
        [(a_var, a_type), (b_var, b_type), (curr_var, curr_type)],
        mk_exists([(res_var, res_type)],
            mk_and([
                mk_primitive("rel_primitive_concat", [res_var, b_var, curr_var]),
                mk_primitive("rel_primitive_concat", [a_var, sep, res_var])
            ])
        )
    )

    output_term = _translate_term(ctx, target)[0]

    return lqp_bundle.ir.Reduce(meta=None, op=op, body=body, terms=[output_term])

def _translate_var(ctx: TranslationCtx, term: ir.Var) -> lqp_bundle.ir.Var:
    name = ctx.var_names.get_name_by_id(term.id, term.name)
    return lqp_bundle.ir.Var(name=name, meta=None)
