from relationalai.early_access.lqp import ir as lqp
from relationalai.early_access.lqp.utils import UniqueNames

rel_to_lqp = {
    "+": "rel_primitive_add",
    "-": "rel_primitive_subtract",
    "*": "rel_primitive_multiply",
    "/": "rel_primitive_divide",
    "=": "rel_primitive_eq",
    "!=": "rel_primitive_neq",
    "<=": "rel_primitive_lt_eq",
    ">=": "rel_primitive_gt_eq",
    ">": "rel_primitive_gt",
    "<": "rel_primitive_lt",
    "abs": "rel_primitive_abs",
    "construct_date": "rel_primitive_construct_date",
    "construct_datetime": "rel_primitive_construct_datetime",
    "construct_datetime_ms_tz": "rel_primitive_construct_datetime",
    "parse_date": "rel_primitive_parse_date",
    "parse_datetime": "rel_primitive_parse_datetime",
    "parse_decimal": "rel_primitive_parse_decimal",
    "parse_decimal64": "rel_primitive_parse_decimal",
    "parse_decimal128": "rel_primitive_parse_decimal",
    "starts_with": "rel_primitive_starts_with",
    "ends_with": "rel_primitive_ends_with",
    "contains": "rel_primitive_contains",
    "num_chars": "rel_primitive_num_chars",
    "substring": "rel_primitive_substring",
    "like_match": "rel_primitive_like_match",
    "concat": "rel_primitive_concat",
    "replace": "rel_primitive_replace",
    "date_year": "rel_primitive_date_year",
    "date_month": "rel_primitive_date_month",
    "date_day": "rel_primitive_date_day",
    "range": "rel_primitive_range",
}

agg_to_lqp = {
    "min": "rel_primitive_min",
    "max": "rel_primitive_max",
    "sum": "rel_primitive_add",
    "count": "rel_primitive_add", # count is a sum of 1s
    "rel_primitive_solverlib_ho_appl": "rel_primitive_solverlib_ho_appl",
}

def relname_to_lqp_name(name: str) -> str:
    # TODO: do these proprly
    if name in rel_to_lqp:
        return rel_to_lqp[name]
    else:
        # If we don't have a mapping for the built-in, we just pass it through as-is.
        return name

# We take the name and type of the variable that we're summing over, so that we can generate
# recognizable names for the variables in the reduce operation and preserve the type.
def lqp_avg_op(names: UniqueNames, op_name: str, sum_name: str, sum_type: lqp.RelType) -> lqp.Abstraction:
    count_type = lqp.PrimitiveType.INT
    vars = [
        (lqp.Var(name=names.get_name(sum_name), meta=None), sum_type),
        (lqp.Var(name=names.get_name("counter"), meta=None), count_type),
        (lqp.Var(name=names.get_name(sum_name), meta=None), sum_type),
        (lqp.Var(name=names.get_name("one"), meta=None), count_type),
        (lqp.Var(name=names.get_name("sum"), meta=None), sum_type),
        (lqp.Var(name=names.get_name("count"), meta=None), count_type),
    ]

    x1 = vars[0][0]
    x2 = vars[1][0]
    y1 = vars[2][0]
    y2 = vars[3][0]
    sum = vars[4][0]
    count = vars[5][0]

    body = lqp.Conjunction(
        args=[
            lqp.Primitive(name="rel_primitive_add", terms=[x1, y1, sum], meta=None),
            lqp.Primitive(name="rel_primitive_add", terms=[x2, y2, count], meta=None)
        ],
        meta=None
    )
    return lqp.Abstraction(vars=vars, value=body, meta=None)

# Default handler for aggregation operations in LQP.
def lqp_agg_op(names: UniqueNames, op_name: str, aggr_arg_name: str, aggr_arg_type: lqp.RelType) -> lqp.Abstraction:
    x = lqp.Var(name=names.get_name(f"x_{aggr_arg_name}"), meta=None)
    y = lqp.Var(name=names.get_name(f"y_{aggr_arg_name}"), meta=None)
    z = lqp.Var(name=names.get_name(f"z_{aggr_arg_name}"), meta=None)
    ts = [(x, aggr_arg_type), (y, aggr_arg_type), (z, aggr_arg_type)]

    name = agg_to_lqp.get(op_name, op_name)
    body = lqp.Primitive(name=name, terms=[x, y, z], meta=None)

    return lqp.Abstraction(vars=ts, value=body, meta=None)

def lqp_operator(names: UniqueNames, op_name: str, aggr_arg_name: str, aggr_arg_type: lqp.RelType) -> lqp.Abstraction:
    # TODO: Can we just pass through unknown operations?
    if op_name not in agg_to_lqp:
        raise NotImplementedError(f"Unsupported aggregation: {op_name}")

    return lqp_agg_op(names, op_name, aggr_arg_name, aggr_arg_type)
