from collections import defaultdict
from copy import deepcopy
from numbers import Number
import textwrap
from typing import Any
from .. import dsl, errors
import gravis as gv
from . import rel, Vars

#--------------------------------------------------
# Errors
#--------------------------------------------------

def invalid_param(param:str, message:str):
    source = errors.Errors.call_source(4)
    marked = errors.mark_source(source, source.line, source.line) if source else ""
    content = f"""
    The [yellow]{param}[/yellow] parameter {message}

    {marked}
    """
    errors.print_source_error(source, "Invalid algorithm parameter", content)
    raise ValueError(message)

#--------------------------------------------------
# Validation
#--------------------------------------------------

def between(param:str,value, min, max,  type_=Number):
    if (not isinstance(value, type_)) or (value < min or value > max):
        invalid_param(param, f"must be a {type_.__name__.lower()} between {min} and {max}, got {value!r}")

def positive(param:str, value, type_=Number):
    if (not isinstance(value, type_)) or (value <= 0):
        invalid_param(param, f"must be a positive {type_.__name__.lower()}, got {value!r}")

def nonnegative(param:str, value, type_=Number):
    if (not isinstance(value, type_)) or (value < 0):
        invalid_param(param, f"must be a non-negative {type_.__name__.lower()}, got {value!r}")

#--------------------------------------------------
# Helpers
#--------------------------------------------------

def unwrap(v):
    if v is None:
        return None
    return rel.pyrel_default(rel.pyrel_unwrap, v, v)


#--------------------------------------------------
# Algos
#--------------------------------------------------

class Compute:
    def __init__(self, graph:'Graph'):
        self._graph = graph
        self._lib = dsl.global_ns.graphlib

    def _config(self, **kwargs):
        return dsl.InlineRelation(self._graph.model, [
            *[(dsl.Symbol(k), v) for k, v in kwargs.items()],
        ])

    def _lookup(self, v):
        res = Vars(1)
        lookup_rel = getattr(rel, f"{self._graph._graph_ref()}_lookup")
        rel.pyrel_default(lookup_rel, v, v, res)
        return res

    # --------------------------------------------------
    # Degree
    # --------------------------------------------------

    def degree(self, node, weight=None):
        if not weight:
            return self._lib.degree(self._graph, unwrap(node))
        return self._lib.weighted_degree(self._graph, unwrap(node))

    def min_degree(self):
        return self._lib.min_degree(self._graph)

    def max_degree(self):
        return self._lib.max_degree(self._graph)

    def avg_degree(self):
        return self._lib.average_degree(self._graph)

    def indegree(self, node):
        return self._lib.indegree(self._graph, unwrap(node))

    def min_indegree(self):
        return self._lib.min_indegree(self._graph)

    def max_indegree(self):
        return self._lib.max_indegree(self._graph)

    def avg_indegree(self):
        return self._lib.average_indegree(self._graph)

    def outdegree(self, node):
        return self._lib.outdegree(self._graph, unwrap(node))

    def min_outdegree(self):
        return self._lib.min_outdegree(self._graph)

    def max_outdegree(self):
        return self._lib.max_outdegree(self._graph)

    def avg_outdegree(self):
        return self._lib.average_outdegree(self._graph)

    # --------------------------------------------------
    # Basics
    # --------------------------------------------------

    def num_nodes(self):
        return self._lib.num_nodes(self._graph)

    def num_edges(self):
        return self._lib.num_edges(self._graph)

    # --------------------------------------------------
    # Distance
    # --------------------------------------------------

    def distance(self, node1, node2):
        return self._lib.distance(self._graph, unwrap(node1), unwrap(node2))

    def weighted_distance(self, node1, node2):
        return self._lib.weighted_distance(self._graph, unwrap(node1), unwrap(node2))

    def diameter_range(self):
        min, max = Vars(2)
        self._lib.diameter_range(self._graph, "min", min)
        self._lib.diameter_range(self._graph, "max", max)
        return (min, max)

    # --------------------------------------------------
    # Connectivity
    # --------------------------------------------------

    def is_connected(self):
        model = self._graph.model
        with model.ordered_choice() as v:
            with model.scope():
                self._lib.is_connected(self._graph)
                v.add(True)
            with model.scope():
                v.add(False)
        return v

    def reachable_from(self, node: dsl.Producer) -> dsl.Expression:
        return self._lib.transitive_closure(self._graph, unwrap(node))

    def is_reachable(self, node1: dsl.Producer, node2: dsl.Producer) -> dsl.Expression:
        model = self._graph.model
        with model.ordered_choice() as v:
            with model.scope():
                self._lib.transitive_closure(self._graph, unwrap(node1), unwrap(node2))
                v.add(True)
            with model.scope():
                v.add(False)
        return v

    # --------------------------------------------------
    # Triangles
    # --------------------------------------------------

    def num_triangles(self, node=None):
        if node is None:
            return self._lib.num_triangles(self._graph)
        return self._lib.triangle_count(self._graph, unwrap(node))

    def is_triangle(self, node1, node2, node3):
        model = dsl.get_graph()
        with model.ordered_choice() as v:
            with model.scope():
                self._lib.triangle(self._graph, unwrap(node1), unwrap(node2), unwrap(node3))
                v.add(True)
            with model.scope():
                v.add(False)
        return v

    def triangles(self, node=None):
        # If no node is provided, return all triangles in the graph,
        # unique up to ordering of the nodes.
        if node is None:
            x, y, z = Vars(3)
            self._lib.unique_triangle(self._graph, x, y, z)
            lx = self._lookup(x)
            ly = self._lookup(y)
            lz = self._lookup(z)
            return (lx, ly, lz)
        # If a node is provided, return all triangles that include that node,
        # unique up to order of the nodes.
        x, y = Vars(2)
        model = dsl.get_graph()
        node = unwrap(node)
        with model.union(dynamic=True) as t:
            for (a, b, c) in ([node, x, y], [x, node, y], [x, y, node]):
                with model.scope():
                    self._lib.unique_triangle(self._graph, a, b, c)
                    t.add(node, a=a, b=b, c=c)
        la = self._lookup(t.a)
        lb = self._lookup(t.b)
        lc = self._lookup(t.c)
        return (la, lb, lc)

    # --------------------------------------------------
    # Clustering
    # --------------------------------------------------

    def local_clustering_coefficient(self, node):
        if not self._graph.undirected:
            invalid_param("graph", "must be undirected")
        return self._lib.local_clustering_coefficient(self._graph, unwrap(node))

    def avg_clustering_coefficient(self):
        if not self._graph.undirected:
            invalid_param("graph", "must be undirected")
        return self._lib.average_clustering_coefficient(self._graph)

    # --------------------------------------------------
    # Link Prediction
    # --------------------------------------------------

    def adamic_adar(self, node1, node2):
        return self._lib.adamic_adar(self._graph, unwrap(node1), unwrap(node2))

    def preferential_attachment(self, node1, node2):
        return self._lib.preferential_attachment(self._graph, unwrap(node1), unwrap(node2))

    def common_neighbor(self, node1, node2):
        return self._lib.common_neighbor(self._graph, unwrap(node1), unwrap(node2))

    # --------------------------------------------------
    # Similarity
    # --------------------------------------------------

    def cosine_similarity(self, node1, node2):
        if not self._graph.undirected:
            invalid_param("graph", "must be undirected")
        return self._lib.cosine_similarity(self._graph, unwrap(node1), unwrap(node2))

    def weighted_cosine_similarity(self, node1, node2):
        if not self._graph.undirected:
            invalid_param("graph", "must be undirected")
        return self._lib.weighted_cosine_similarity(self._graph, unwrap(node1), unwrap(node2))

    def jaccard_similarity(self, node1, node2):
        return self._lib.jaccard_similarity(self._graph, unwrap(node1), unwrap(node2))

    def weighted_jaccard_similarity(self, node1, node2):
        return self._lib.weighted_jaccard_similarity(self._graph, unwrap(node1), unwrap(node2))

    # --------------------------------------------------
    # Centrality
    # --------------------------------------------------

    def pagerank(self, node, damping_factor=0.85, tolerance=1e-6, max_iter=20):
        between("damping_factor", damping_factor, 0, 1)
        positive("tolerance", tolerance)
        positive("max_iter", max_iter, type_=int)

        config = self._config(
            graph=self._graph,
            damping_factor=float(damping_factor),
            tolerance=float(tolerance),
            max_iter=max_iter,
        )
        return self._lib.pagerank(config, unwrap(node))

    def betweenness_centrality(self, node):
        return self._lib.betweenness_centrality(self._graph, unwrap(node))

    def degree_centrality(self, node):
        return self._lib.degree_centrality(self._graph, unwrap(node))

    def eigenvector_centrality(self, node):
        if not self._graph.undirected:
            invalid_param("graph", "must be undirected")
        return self._lib.eigenvector_centrality(self._graph, unwrap(node))

    def weighted_degree_centrality(self, node):
        return self._lib.weighted_degree_centrality(self._graph, unwrap(node))

    # --------------------------------------------------
    # Community Detection
    # --------------------------------------------------

    def weakly_connected_component(self, node):
        component = self._lib.weakly_connected_component(self._graph, unwrap(node))
        return self._lookup(component)

    def triangle_community(self, node):
        return self._lib.triangle_community(self._graph, unwrap(node))

    def infomap(
            self,
            node,
            max_levels: int = 4,
            max_sweeps: int = 8,
            level_tolerance: float = 0.01,
            sweep_tolerance: float = 0.0001,
            teleportation_rate: float = 0.15,
            visit_rate_tolerance: float = 1e-15,
            randomization_seed: int | None = None,
    ):
        positive("max_levels", max_levels, type_=int)
        nonnegative("max_sweeps", max_sweeps, type_=int)
        nonnegative("level_tolerance", level_tolerance)
        nonnegative("sweep_tolerance", sweep_tolerance)
        positive("teleportation_rate", teleportation_rate)
        between("teleportation_rate", teleportation_rate, 0, 1)
        positive("visit_rate_tolerance", visit_rate_tolerance)
        _config_dict = {
            "graph": self._graph,
            "max_levels": max_levels,
            "max_sweeps": max_sweeps,
            "level_tolerance": float(level_tolerance),
            "sweep_tolerance": float(sweep_tolerance),
            "teleportation_rate": float(teleportation_rate),
            "visit_rate_tolerance": float(visit_rate_tolerance),
        }
        if randomization_seed is not None:
            positive("randomization_seed", randomization_seed, type_=int)
            _config_dict["randomization_seed"] = randomization_seed

        config = self._config(**_config_dict)
        return self._lib.infomap(config, unwrap(node))

    def louvain(
            self,
            node,
            max_levels: int = 4,
            max_sweeps: int = 8,
            level_tolerance: float = 0.01,
            sweep_tolerance: float = 0.0001,
            randomization_seed: int | None = None,
    ):
        positive("max_levels", max_levels, type_=int)
        nonnegative("max_sweeps", max_sweeps, type_=int)
        nonnegative("level_tolerance", level_tolerance)
        nonnegative("sweep_tolerance", sweep_tolerance)
        _config_dict = {
            "graph": self._graph,
            "max_levels": max_levels,
            "max_sweeps": max_sweeps,
            "level_tolerance": float(level_tolerance),
            "sweep_tolerance": float(sweep_tolerance),
        }
        if randomization_seed is not None:
            positive("randomization_seed", randomization_seed, type_=int)
            _config_dict["randomization_seed"] = randomization_seed
        config = self._config(**_config_dict)
        return self._lib.louvain(config, unwrap(node))

    def label_propagation(self, node, max_sweeps: int = 1000, randomization_seed: int | None = None):
        positive("max_sweeps", max_sweeps, type_=int)
        _config_dict = {
            "graph": self._graph,
            "max_sweeps": max_sweeps,
        }
        if randomization_seed is not None:
            positive("randomization_seed", randomization_seed, type_=int)
            _config_dict["randomization_seed"] = randomization_seed
        config = self._config(**_config_dict)
        return self._lib.label_propagation(config, unwrap(node))

#--------------------------------------------------
# Edge
#--------------------------------------------------

class EdgeInstance:
    def __init__(self, edge:'Edge', from_:Any, to:Any, kwargs:dict={}):
        self._edge = edge
        self._graph = edge._graph
        self.from_ = from_
        self.to = to
        for k, v in kwargs.items():
            self._edge._prop(k)(from_, to, v)

    def __getattr__(self, name:str):
        v = Vars(1)
        self._edge._prop(name)(self.from_, self.to, v)
        return v

    def set(self, **kwargs):
        for k, v in kwargs.items():
            self._edge._prop(k).add(self.from_, self.to, v)
        return self

    def _to_var(self):
        raise Exception("Edges can't be returned directly, you can return the from_ and to properties individually")

class Edge:
    def __init__(self, graph:'Graph'):
        self._graph = graph
        self._type = dsl.RawRelation(self._graph.model, f"graph{self._graph.id}_edges", 2)
        self._props = {}

    def _prop(self, name:str):
        if name not in self._props:
            self._props[name] = dsl.RawRelation(self._graph.model, f"graph{self._graph.id}_edge_{name}", 3)
        return self._props[name]

    def _is_weighted(self):
        return "weight" in self._props

    def extend(self, prop:'dsl.Property', **kwargs):
        type = prop._provider
        with self._graph.model.rule():
            t = type()
            self.add(t, getattr(t, prop._name), **kwargs)

    def add(self, from_:Any, to:Any, **kwargs):
        self._type.add(from_, to)
        for k, v in kwargs.items():
            self._prop(k).add(from_, to, v)
        return EdgeInstance(self, from_, to, {})

    def __call__(self, from_:Any=None, to:Any=None, **kwargs):
        if from_ is None:
            from_ = Vars(1)
        if to is None:
            to = Vars(1)
        self._type(from_, to)
        return EdgeInstance(self, from_, to, kwargs)

    def __getattribute__(self, __name: str) -> Any:
        if __name in ["add", "extend"] or __name.startswith("_"):
            return super().__getattribute__(__name)
        return self._props[__name]

#--------------------------------------------------
# Graph
#--------------------------------------------------

class Graph:
    def __init__(self, model:dsl.Graph, undirected=False, weighted=False, default_weight=1.0, with_isolated_nodes=False):
        self.model = model
        self.id = dsl.next_id()
        self.compute = Compute(self)
        self.Node = dsl.Type(model, "nodes", scope=f"graph{self.id}_")
        self.Edge = Edge(self)
        self.undirected = undirected
        self.weighted = weighted
        self.default_weight=default_weight
        self._last_fetch = None

        graph_type = "::graphlib::undirected_graph" if undirected else "::graphlib::directed_graph"
        if weighted:
            unwrapped = f"""
            def {self.Edge._prop("weight")._name}_unwrapped(au, bu, w): {{
                exists((a, b) |
                    {self.Edge._prop("weight")._name}(a, b, w) and
                    pyrel_default(pyrel_unwrap, a, a, au) and
                    pyrel_default(pyrel_unwrap, b, b, bu)
                )
            }}
            """
            create_graph = f"""{graph_type}[{self.Edge._prop("weight")._name}_unwrapped]"""
        else:
            unwrapped = f"""
            def {self.Edge._type._name}_unwrapped(au, bu): {{
                exists((a, b) |
                    {self.Edge._type._name}(a, b) and
                    pyrel_default(pyrel_unwrap, a, a, au) and
                    pyrel_default(pyrel_unwrap, b, b, bu)
                )
            }}
            """
            create_graph = f"{graph_type}[{self.Edge._type._name}_unwrapped]"

        node_def = ""
        if with_isolated_nodes:
            node_def = f"def {self._graph_ref()}[:node]: {{pyrel_unwrap[{self.Node._type.name}]}}"

        self.model.install_raw(textwrap.dedent(f"""
        declare {self.Node._type.name}
        declare {self.Edge._type._name}
        {unwrapped}
        @inline
        def {self._graph_ref()}_lookup(uw, orig): {{ {self.Node._type.name}(orig) and pyrel_unwrap(orig, uw) }}

        {f"declare {self.Edge._prop('weight')._name}" if weighted else ""}
        def {self._graph_ref()} {{{create_graph}}}
        {node_def}
        """))

        # Add a rule that makes all nodes used in edges are also added to
        # the nodes relation
        with model.rule():
            a, b = dsl.Vars(2)
            self.Edge._type(a, b)
            self.Node.add(a)
            self.Node.add(b)

    def _graph_ref(self):
        return f"graph{self.id}"

    def _to_var(self):
        return getattr(rel, self._graph_ref())._to_var()

    def _is_weighted(self):
        return self.weighted
        # return self.edges._is_weighted()

    #--------------------------------------------------
    # Fetch
    #--------------------------------------------------

    def fetch(self):
        code = []
        code.append(f"def output(:nodes, n): {self.Node._type.name}(n)")
        for prop in self.Node._type.properties:
            scope = self.Node._scope
            code.append(f"def output(:nodes, :{prop.name.removeprefix(scope)}, n, v): {prop.name}(n, v)")
        code.append(f"def output(:edges, a,b): {self.Edge._type._name}(a,b)")
        for name, prop in self.Edge._props.items():
            code.append(f"def output(:edges, :{name}, a, b, v): {prop._name}(a, b, v)")

        output = {"nodes": defaultdict(dict), "edges": defaultdict(dict)}
        results = self.model.exec_raw("\n".join(code), raw_results=True)
        for set_ in results.results:
            path = [v[1:] for v in set_["relationId"].split("/")[2:] if v[0] == ":"]
            if path[0] not in ["nodes", "edges"]:
                continue
            cur = output[path[0]]
            if path[0] == "nodes":
                if len(path) == 1:
                    for (n,) in set_["table"].itertuples(index=False):
                        cur[n]
                else:
                    for (n, v) in set_["table"].itertuples(index=False):
                        cur[n][path[1]] = v
            elif path[0] == "edges":
                if len(path) == 1:
                    for (a, b) in set_["table"].itertuples(index=False):
                        # If the graph is undirected, normalize the order of
                        # the source and target nodes.
                        if self.undirected and a > b:
                            a, b = b, a
                        cur[(a, b)]
                else:
                    for (a, b, v) in set_["table"].itertuples(index=False):
                        if self.undirected and a > b:
                            a, b = b, a
                        cur[(a, b)][path[1]] = v
            else:
                raise Exception(f"Unexpected path: {path}")
        self._last_fetch = output
        return output

    #--------------------------------------------------
    # Visualize
    #--------------------------------------------------

    default_visual_props = {
        "node": {
            "color": "black",
            "opacity": 1,
            "size": 10,
            "shape": "circle",
            "border_color": "black",
            "border_size": 1,
            "label_color": "black",
            "label_size": 10,
        },
        "edge": {
            "color": "#999",
            "opacity": 1,
            "size": 2,
            "shape": "circle",
            "border_color": "#999",
            "border_size": 1,
            "label_color": "black",
            "label_size": 10,
            "arrow_size": 4,
            "arrow_color": "#999",
        }
    }

    _style_map = {
        "arrow_size": "arrow_size",
        "arrow_color": "arrow_color",
    }

    def _visual_props(self, prop_def, metadata):
        for k, v in prop_def.items():
            if callable(v):
                metadata[k] = v(metadata)
        return metadata

    def _visual_dict(self, style: dict, use_cache = False) -> dict:
        data = self._last_fetch if use_cache else None
        if not data:
            data = self.fetch()
        style = deepcopy(style)
        if "node" not in style:
            style["node"] = {}
        if "edge" not in style:
            style["edge"] = {}
        merged_style = deepcopy(self.default_visual_props)
        for category in ["node", "edge"]:
            for k, v in style.get(category, {}).items():
                if not callable(v):
                    merged_style[category][k] = v
        flat_style = {
            self._style_map.get(k, f"{type}_{k}"): v
            for type, category in merged_style.items()
            for k, v in category.items()
        }

        return {
            "graph": {
                "directed": not self.undirected,
                "metadata": flat_style,
                "nodes": {
                    node_id: {
                        **({"label": info["label"]} if "label" in info else {}),
                        "metadata": self._visual_props(style["node"], info),
                    }
                    for (node_id, info) in data["nodes"].items()
                },
                "edges": [
                    {"source": source, "target": target, "metadata": self._visual_props(style["edge"], info)}
                    for ((source, target), info) in data["edges"].items()
                ],
            }
        }

    def visualize(self, three=False, style: dict | None = None, use_cache = False, **kwargs):
        style = style if style is not None else {"node": {}, "edge": {}}
        vis = gv.vis if not three else gv.three
        graph_dict = self._visual_dict(style=style, use_cache=use_cache)
        # Use defaults for the following kwargs if not provided
        new_kwargs = {
            "node_label_data_source": "label",
            "edge_label_data_source": "label",
            "show_edge_label": True,
            "edge_curvature": 0.4,
        } | kwargs
        fig = vis(graph_dict, **new_kwargs)
        return fig

#--------------------------------------------------
# Path
#--------------------------------------------------

class Path:
    def __init__(self, *args):
        self.edges = []
        self.nodes = []
        pass

    def __getitem__(self, item):
        return self

#--------------------------------------------------
# Exports
#--------------------------------------------------

__all__ = ["Graph", "Path"]
