from __future__ import annotations
from typing import Any, Union
import textwrap

from relationalai.early_access.builder import builder as b, std
from relationalai.early_access.rel.executor import RelExecutor
from .common import make_name
from relationalai.experimental.solvers import Solver

_Any = Union[b.Producer, str, float, int]
_String = Union[b.Producer, str]
_Number = Union[b.Producer, float, int]

class SolverModelDev:
    def __init__(self, model: b.Model, num_type: str):
        self._model = model
        assert num_type in ("cont", "int")
        data_type = "int" if num_type == "int" else "float"
        self._data_type = data_type
        self._variable_relationships = set()
        self._expr_id = 0
        Relationship = model.Relationship

        # TODO workaround for QB not distinguishing string names for Concepts/Relationships, only a problem with raw Rel for result extraction maybe
        sm_id = next(b._global_id)
        def _name(s: str) -> str:
            return f"{s}_{sm_id}"

        # variable relations
        self.variables = Relationship("{node} is variable", short_name=_name("variables"))
        self.variable_type = Relationship("{node} has {type:str}", short_name=_name("variable_type")) # TODO use codes/enum
        self.variable_name = Relationship("{node} has {name:str}", short_name=_name("variable_name"))
        self.variable_start = Relationship(f"{{node}} has {{start:{data_type}}}", short_name=_name("variable_start"))

        # expression relations
        self.operator = Relationship("{node} has {op:str}", short_name=_name("operator")) # TODO use codes/enum
        self.ordered_args_hash = Relationship("{node} arg {i:int} is {arg}", short_name=_name("ordered_args_hash"))
        self.ordered_args_data = Relationship(f"{{node}} arg {{i:int}} is {{arg:{data_type}}}", short_name=_name("ordered_args_data"))
        self.unordered_args_hash = Relationship("{node} arg {i} is {arg}", short_name=_name("unordered_args_hash"))
        self.min_objectives = Relationship("{node} is min objective", short_name=_name("min_objectives"))
        self.max_objectives = Relationship("{node} is max objective", short_name=_name("max_objectives"))
        self.constraints = Relationship("{node} is constraint", short_name=_name("constraints"))
        self.expression_name = Relationship("{node} has {name:str}", short_name=_name("expression_name")) # TODO use codes/enum

        # solver result relations
        self.result_info = model.Relationship("{key:str} has {val:str}", short_name=_name("result_info"))
        self.point = model.Relationship(f"{{var}} has {{point:{data_type}}}", short_name=_name("point"))
        self.points = model.Relationship(f"point {{i:int}} for {{var}} has {{point:{data_type}}}", short_name=_name("points"))

    def solve_for(self, expr: b.Relationship | b.Fragment, populate: bool = True, **kwargs):
        where = []
        if isinstance(expr, b.Fragment):
            assert expr._select and len(expr._select) == 1 and expr._where, "Fragment input for `solve_for` must have exactly one select and a where clause"
            rel = expr._select[0]
            where.extend(expr._where)
        elif isinstance(expr, b.Relationship):
            rel = expr
        else:
            raise ValueError(f"Invalid expression type {type(expr)} for `solve_for`; must be a Relationship or Fragment")
        assert rel._parent and rel._short_name, "Relationship for `solve_for` must have a parent and a short name"
        assert rel not in self._variable_relationships

        self._variable_relationships.add(rel)
        node = b.Hash.ref()
        defs = [self.variables(node)]

        # handle optional variable properties
        new_kwargs = kwargs.copy()
        if "type" not in new_kwargs:
            new_kwargs["type"] = "cont" if self._data_type == "float" else "int"
        for (key, val) in new_kwargs.items():
            if key == "name":
                assert isinstance(val, (_Any, list)), f"Expected {key} to be a value or list, got {type(val)}"
                defs.append(self.variable_name(node, make_name(val)))
            elif key == "type":
                assert val in ("cont", "int", "bin"), f"Unsupported variable type {val} for `solve_for`; must be cont, int, or bin"
                defs.append(self.variable_type(node, val))
            elif key in ("lower", "upper", "fixed", "start"):
                assert isinstance(val, _Number), f"Expected {key} to be a number, got {type(val)}"
                if key == "start":
                    defs.append(self.variable_start(node, val))
                else:
                    # TODO use lower/upper/fixed tables in future rather than making constraints
                    op = ">=" if key == "lower" else ("<=" if key == "upper" else "=")
                    self.satisfy(b.require(b.Expression(b.Relationship.builtins[op], rel, val)).where(*where))
            else:
                raise ValueError(f"Invalid keyword argument {key} for `solve_for`")

        where.append(_make_hash((rel._short_name, rel._parent), node))
        b.define(*defs).where(*where)

        if populate:
            # get variable values from the result point (populated by the solver)
            # TODO maybe instead delete/insert into variable relationships after solve.
            val = b.Number.ref()
            b.define(rel(val)).where(self.point(node, val), *where)

        return None

    def minimize(self, expr: _Number, name: _String | list | None = None):
        assert isinstance(expr, _Number)
        return self._handle_expr(self.min_objectives, expr, name)

    def maximize(self, expr: _Number, name: _String | list | None = None):
        assert isinstance(expr, _Number)
        return self._handle_expr(self.max_objectives, expr, name)

    def satisfy(self, expr: b.Fragment, check: bool = False, name: _String | list | None = None):
        assert expr._require, "Fragment input for `satisfy` must have a require clause"
        assert not expr._select and not expr._define, "Fragment input for `satisfy` must not have a select or define clause"
        if not check:
            # remove the `require` from the model roots so it is not checked
            b._remove_roots([expr])
        for req in expr._require:
            self._handle_expr(self.constraints, req, name, where=expr._where)
        return None

    def _handle_expr(self, root_type, expr, name: _String | list | None, where: list = []):
        ctx = ExprContext(self)
        ctx.where.extend(where)
        root = _rewrite(expr, ctx)
        assert root is not None, f"Cannot symbolify {expr}"
        ctx.define.append(root_type(root))
        # if name is not None:
        #     # TODO not printing expression names yet, so skip for now
        #     assert isinstance(name, (_String, list)), f"Expected name to be a string or list, got {type(name)}"
        #     ctx.define.append(self.expression_name(root, make_name(name)))
        self._handle_ctx(ctx)
        return None

    def _handle_ctx(self, ctx: ExprContext, where: list = []):
        all_where = ctx.where + where
        b.define(*ctx.define).where(*all_where)
        for subctx in ctx.subcontext:
            self._handle_ctx(subctx, where=all_where)

    # print the variables and components of the model in human-readable format
    def print(self, verbose: bool = True, expr_names: bool = False):
        # summarize variables
        var_types = ["int", "bin", "cont"]
        vi = b.Hash.ref()
        var_counts = b.select(*[b.count(vi).where(self.variable_type(vi, t)) | 0 for t in var_types]).to_df()
        assert var_counts.shape == (1, 3)
        (int_count, bin_count, cont_count) = var_counts.iloc[0]

        # summarize expressions
        expr_types = [
            (self.min_objectives, "Minimization objectives"),
            (self.max_objectives, "Maximization objectives"),
            (self.constraints, "Constraints"),
        ]
        ei = b.Hash.ref()
        expr_counts = b.select(*[b.count(ei).where(E(ei)) | 0 for (E, _) in expr_types]).to_df()
        assert expr_counts.shape == (1, 3)
        (min_count, max_count, cons_count) = expr_counts.iloc[0]

        # print summary
        print("Solver model has:")
        print(f"• {int_count} integer variables, {bin_count} binary variables, {cont_count} continuous variables")
        print(f"• {min_count} minimization objectives, {max_count} maximization objectives, {cons_count} constraints")
        if not verbose:
            return None

        # print variable names
        vn, vt = b.String.ref(), b.String.ref()
        var_names = b.select(vn.alias("name"), vt.alias("type")).where(self.variable_name(vi, vn), self.variable_type(vi, vt)).to_df()
        assert not var_names.empty, "No variable names found in the model"
        print("Variables:")
        print(var_names.to_string(index=False, header=True))

        # print expressions TODO print expression names optionally
        for (E, s) in expr_types:
            expr_strings = self._expr_strings(E)
            if expr_strings:
                print(s + ":")
                print("\n".join(expr_strings))
        return None

    def _expr_strings(self, roots):
        # materialize data
        iv = b.Hash.ref()
        roots_df = b.select(iv).where(roots(iv)).to_df()
        if roots_df.empty:
            return []
        nodes = roots_df.iloc[:, 0].tolist()

        sv = b.String.ref()
        names_dict = {k: v for k, v in b.select(iv, sv).where(self.variable_name(iv, sv))}
        ops_dict = {k: v for k, v in b.select(iv, sv).where(self.operator(iv, sv))}

        args_dict = {}
        iv2 = b.Hash.ref()
        iv3 = b.Hash.ref()
        # unordered args
        for k, _, v in b.select(iv, iv2, iv3).where(self.unordered_args_hash(iv, iv2, iv3)):
            if k not in args_dict:
                args_dict[k] = []
            args_dict[k].append(v)
        ordered_args_dict = {}
        iv2 = b.Integer.ref()
        for k, i, v in b.select(iv, iv2, iv3).where(self.ordered_args_hash(iv, iv2, iv3)):
            if k not in ordered_args_dict:
                ordered_args_dict[k] = []
            ordered_args_dict[k].append((i, v))
        vv = b.python_types_str_to_concepts[self._data_type].ref()
        for k, i, v in b.select(iv, iv2, vv).where(self.ordered_args_data(iv, iv2, vv)):
            if k not in ordered_args_dict:
                ordered_args_dict[k] = []
            ordered_args_dict[k].append((i, v))
        # convert ordered args to dict
        for k, t in ordered_args_dict.items():
            # sort by index
            t.sort(key=lambda x: x[0])
            args_dict[k] = [v for (_, v) in t]

        # compute strings
        return sorted(_expr_strings_rec(n, names_dict, ops_dict, args_dict) for n in nodes)

    # solve the model given a solver and solver options
    # TODO for now we are materializing the expression graph as dataframes and writing to CSVs in python on the client side, then sending the CSVs along with the solver options as HTTP parameters, skipping the export of a model file from the engine.
    def solve(self, solver: Solver, log_to_console: bool = False, **kwargs):
        # validate options
        for k, v in kwargs.items():
            if not isinstance(k, str):
                raise ValueError(f"Invalid parameter key. Expected string, got {type(k)} for {k}.")
            if not isinstance(v, (int, float, str, bool)):
                raise ValueError(
                    f"Invalid parameter value. Expected string, integer, float, or boolean, got {type(v)} for {k}."
                )

        # 1. Query the model relations and convert to CSV strings.
        print("query model relations")
        def _csv(t):
            return t.to_df().to_csv(index=False, header=False)
        model_relations: dict[str, Any] = {
            "variable_type": _csv(self.variable_type),
            "variable_name": _csv(self.variable_name),
            "operator": _csv(self.operator),
            "ordered_args_hash": _csv(self.ordered_args_hash),
            "ordered_args_data": _csv(self.ordered_args_data),
            "unordered_args_hash": _csv(self.unordered_args_hash),
            "min_objectives": _csv(self.min_objectives),
            "max_objectives": _csv(self.max_objectives),
            "constraints": _csv(self.constraints),
        }

        # 2. Execute job and wait for completion.
        print("execute solver job")
        kwargs["model_relations"] = model_relations
        payload = {"solver": solver.solver_name.lower(), "options": kwargs}
        job_id = solver._exec_job(payload, log_to_console=log_to_console)

        # 3. Extract result.
        print("extract result")
        executor = self._model._to_executor()
        assert isinstance(executor, RelExecutor)
        executor.execute_raw(textwrap.dedent(f"""
        def raw_result {{
            load_binary["snowflake://APP_STATE.RAI_INTERNAL_STAGE/job-results/{job_id}/result.binpb"]
        }}

        ic result_not_empty("Solver result is empty.") requires not empty(raw_result)

        @no_diagnostics(:EXPERIMENTAL)
        def extracted {{ rel_primitive_solverlib_extract[raw_result] }}

        def delete[:{self.result_info._name}]: {self.result_info._name}
        def delete[:{self.point._name}]: {self.point._name}
        def delete[:{self.points._name}]: {self.points._name}

        def insert(:{self.result_info._name}, key, val):
            exists((k) | string(extracted[k], val) and ::std::mirror::lower(k, key))
        def insert(:{self.point._name}, var, val): extracted(:point, var, val)
        def insert(:{self.points._name}, i, var, val): extracted(:points, var, i, val)
        """), readonly=False)

        print("finished solve")
        return None

    # load a particular point index from `points` into `point`
    # so it is accessible from the variable relationship
    def load_point(self, i: int):
        if not isinstance(i, int) and i >= 0:
            raise ValueError(f"Expected nonnegative integer index for point, got {i}")
        executor = self._model._to_executor()
        assert isinstance(executor, RelExecutor)
        executor.execute_raw(textwrap.dedent(f"""
        def delete[:{self.point._name}]: {self.point._name}
        def insert(:{self.point._name}, var, val): {self.points._name}(int128[{i}], var, val)
        """), readonly=False)
        return None

    # print summary of the solver result
    def summarize_result(self):
        to_get = ["error", "termination_status", "solve_time_sec", "objective_value", "solver_version", "result_count"]
        k, v = b.String.ref(), b.String.ref()
        df = b.select(k, v).where(self.result_info(k, v), k.in_(to_get)).to_df()
        print(df.to_string(index=False, header=False))
        return df

    # select variable names and values in the primal result point(s)
    def variable_values(self, multiple: bool = False):
        var = b.Hash.ref()
        val = b.Number.ref()
        s = b.String.ref()
        if multiple:
            i = b.Integer.ref()
            return b.select(i, s, val).where(self.points(i, var, val), self.variable_name(var, s))
        else:
            return b.select(s, val).where(self.point(var, val), self.variable_name(var, s))

    # get scalar information
    def __getattr__(self, name: str):
        df = None
        # model info
        if name in {"num_variables", "num_constraints", "num_min_objectives", "num_max_objectives"}:
            map = {
                "num_variables": self.variables,
                "num_constraints": self.constraints,
                "num_min_objectives": self.min_objectives,
                "num_max_objectives": self.max_objectives,
            }
            node = b.Hash.ref()
            df = b.select(b.count(node).where(map[name](node)) | 0).to_df()
        # result info
        if name in {"error", "termination_status", "solver_version", "printed_model", "solve_time_sec", "objective_value", "result_count"}:
            val = b.String.ref()
            df = b.select(val).where(self.result_info(name, val)).to_df()
        if df is not None:
            if not df.shape == (1, 1):
                raise ValueError(f"Expected exactly one value for {name}, but df has shape {df.shape}")
            v = df.iloc[0, 0]
            if isinstance(v, str):
                if name == "solve_time_sec":
                    return float(v)
                elif name == "objective_value":
                    return int(v) if self._data_type == "int" else float(v)
                elif name == "result_count":
                    return int(v)
            return v
        return None


class ExprContext():
    def __init__(self, solver_model):
        self.solver_model = solver_model
        self.define = []
        self.where = []
        self.subcontext = []

def _rewrite(expr: b.Producer | b.Fragment, ctx: ExprContext):
    if isinstance(expr, (int, float, str)):
        return None

    elif isinstance(expr, b.Concept):
        return None

    elif isinstance(expr, b.Ref):
        thing = _rewrite(expr._thing, ctx)
        if thing:
            assert isinstance(thing, b.Producer)
            return thing.ref()
        return None

    elif isinstance(expr, (b.Relationship, b.RelationshipRef, b.RelationshipFieldRef)):
        rel = expr if isinstance(expr, b.Relationship) else expr._relationship
        if rel in ctx.solver_model._variable_relationships:
            return std.hash(rel._short_name, expr._parent)
        return None

    elif isinstance(expr, b.Union):
        # return union of the symbolified expressions, if any are symbolic
        args_rewritten = False
        args = []
        for arg in expr._args:
            ra = _rewrite(arg, ctx)
            if ra:
                args_rewritten = True
                args.append(ra)
            else:
                args.append(arg)
        if args_rewritten:
            return b.union(*args)
        return None

    elif isinstance(expr, b.Expression):
        assert isinstance(expr._op, b.Relationship)
        assert _rewrite(expr._op, ctx) is None # shouldn't contain variables
        sym_args = [_rewrite(a, ctx) for a in expr._params]
        if all(a is None for a in sym_args):
            return None
        op = expr._op._name
        assert op in infixs or op in prefix_ops, f"Operator {op} not supported in first order solver rewrites"

        # some arguments involve solver variables, so rewrite the expression
        new_args = sym_args if op in infix_comps else sym_args[:-1] # drop last arg introduced by builder
        sm = ctx.solver_model
        node = b.Hash.ref()
        sm._expr_id += 1
        ctx.define.append(sm.operator(node, op))
        hash_args = []
        for (i, a) in enumerate(new_args):
            if a is None:
                d = sm.ordered_args_data(node, i, expr._params[i])
                hash_args.append(expr._params[i])
            else:
                d = sm.ordered_args_hash(node, i, a)
                hash_args.append(a)
            ctx.define.append(d)
        ctx.where.append(_make_hash((sm._expr_id, *hash_args), node))
        return node

    elif isinstance(expr, b.Aggregate):
        # only the last argument can be symbolic
        pre_args = expr._args[:-1]
        sym_arg = _rewrite(expr._args[-1], ctx)
        if not sym_arg:
            return None
        op = expr._op._name
        assert op in agg_ops, f"Operator {op} not supported in aggregate rewrites"

        sm = ctx.solver_model
        node = b.Hash.ref()
        ctx.where.append(_make_hash((sm._expr_id, *expr._group), node))
        # only apply where conditions in subcontext
        subctx = ExprContext(sm)
        ctx.subcontext.append(subctx)
        subctx.where.extend(expr._where._where)
        arg_hash = b.Hash.ref()
        subctx.where.append(_make_hash((sm._expr_id, *pre_args), arg_hash)) # TODO also add sym_arg here?
        sm._expr_id += 1
        subctx.define.append(sm.operator(node, op))
        subctx.define.append(sm.unordered_args_hash(node, arg_hash, sym_arg)) # TODO what if some values are data not hashes?
        return node

    elif isinstance(expr, b.Fragment):
        # only support selects with one item
        assert not expr._define and not expr._require and len(expr._select) == 1
        sym_select = _rewrite(expr._select[0], ctx)
        if sym_select:
            return b.select(sym_select).where(*expr._where)
        return None

    raise NotImplementedError(f"Solver rewrites cannot handle {expr} of type {type(expr)}")

def _expr_strings_rec(x, names_dict, ops_dict, args_dict):
    if x in names_dict:
        # this is a variable node
        return names_dict[x]
    elif x not in ops_dict:
        # this is a value
        return str(x)

    # this is an expression
    assert x in ops_dict
    arg_strs = []
    op = ops_dict[x]

    for k in args_dict.get(x, []):
        s = _expr_strings_rec(k, names_dict, ops_dict, args_dict)
        if op in infix_ops and k in ops_dict and ops_dict[k] in infixs:
            # add parentheses to avoid precedence issues
            s = f"({s})"
        arg_strs.append(s)

    if op in agg_ops:
        # sort unordered args to improve determinism
        arg_strs.sort()

    if op in infixs:
        # infix operator, e.g. x + y
        assert len(arg_strs) == 2 # TODO could relax e.g. for +, *
        return f"{arg_strs[0]} {op} {arg_strs[1]}"
    else:
        # prefix operator, e.g. abs(x)
        assert len(arg_strs) >= 1
        return f"{op}({', '.join(arg_strs)})"

def _make_hash(tup, res):
    return b.Expression(b.Relationship.builtins["hash"], b.TupleArg(tup), res)

infix_ops = set(["+", "-", "*", "/", "^"])
infix_comps = set(["=", "!=", "<", "<=", ">", ">=", "implies"])
infixs = infix_ops.union(infix_comps)
prefix_ops = set(["abs", "log", "exp"])
agg_ops = set(["sum", "count", "min", "max", "all_different"])

# _variable_types = {
#     "cont": 40,
#     "int": 41,
#     "bin": 42,
# }

# _fo_operators = {
#     "+": 10,
#     "-": 11,
#     "*": 12,
#     "/": 13,
#     "^": 14,
#     "abs": 20,
#     "exp": 21,
#     "log": 22,
#     "range": 50,
# }

# _fo_comparisons = {
#     "=": 30,
#     "!=": 31,
#     "<=": 32,
#     ">=": 33,
#     "<": 34,
#     ">": 35,
#     "implies": 62,
# }

# _ho_operators = {
#     "sum": 80,
#     # "product":81,
#     "min": 82,
#     "max": 83,
#     "count": 84,
#     "all_different": 90,
# }
