import pandas as pd
import numpy as np
import networkx as nx


class SGraph:
    def __init__(self, gr):
        self.gr = gr

    def __copy__(self):
        return SGraph(self.gr.copy())

    def directed(self):
        return nx.is_directed(self.get_nx_graph())

    def clean(self):
        self.gr.clear()

    def subgraph(self, nodes):
        return SGraph((self.gr.subgraph(nodes)).copy())

    def get_ids(self, stable=False):
        if stable:
            return list(self.gr.nodes())
        else:
            return self.gr.nodes()

    def set_attr(self, ind, attr, value):
        self.gr.nodes[ind][attr] = convert(value)

    def set_attrs(self, attr, values):
        nx.set_node_attributes(self.gr, name=attr, values=values)

    def get_attr(self, ind, attr):
        if hasattr(ind, "__iter__") and not isinstance(ind, str):
            res = []
            for i in ind:
                try:
                    res.append(self.gr.nodes[i][attr])
                except:
                    res.append(0)
            return np.array(res)

        try:
            return self.gr.nodes[ind][attr]
        except:
            return 0

    def get_node_data(self, ind):
        return self.gr.nodes[ind]

    def add_node(self, id_n):
        self.gr.add_node(id_n)

    def add_nodes(self, ids):
        self.gr.add_nodes_from(ids)

    def del_node(self, id_n):
        self.gr.remove_node(id_n)

    def del_nodes(self, ids):
        self.gr.remove_nodes_from(ids)

    def copy(self):
        return SGraph(self.gr.copy())

    def get_attributes(self, attr, nodes=None):
        if nodes is None:
            nodes = self.gr.nodes()

        res = []
        for x in nodes:
            try:
                res.append(self.gr.nodes[x][attr])
            except:
                res.append(-1)
        return np.array(res)

    def get_attrs_dict(self, attr):
        return nx.get_node_attributes(self.gr, attr)

    def get_nx_graph(self):
        return self.gr

    def add_edge(self, u, v):
        self.gr.add_edge(u, v)

    def add_edges(self, edges):
        self.gr.add_edges_from(edges)

    def get_edge_data(self, u, v, frm):
        return self.gr[u][v][frm]

    def set_edge_data(self, u, v, to, value):
        self.gr[u][v][to] = value

    def del_edge(self, u, v):
        self.gr.remove_edge(u, v)

    def del_edges_(self, edges):
        self.gr.remove_edges_from(edges)

    def add_path(self, path):
        self.gr.add_path(path)

    def free_in_nodes(self):
        return np.array([n for n in self.get_ids() if self.count_in_degree(n) < self.get_attr(n, "ind")])

    def size(self):
        return self.gr.number_of_nodes()

    def count_in_degree(self, ind, w=None, inn=None):
        res = self.get_in_degrees(ind, w=w)
        if inn is None:
            return res.size
        else:
            return res[np.isin(res, inn)].size

    def count_out_degree(self, out, w=None, inn=None):
        res = self.get_out_degrees(out, w=w)
        if inn is None:
            return res.size
        else:
            return res[np.isin(res, inn)].size

    def count_degree(self, ind, w=None, inn=None):
        res = self.get_degrees(ind, w=w)
        if inn is None:
            return res.size
        else:
            return res[np.isin(res, inn)].size

    def count_edges(self, i, j):
        if i.size > j.size:
            return np.array([x[0] for x in self.gr.in_edges(j) if x[0] in i]).size
        else:
            return np.array([x[1] for x in self.gr.out_edges(i) if x[1] in j]).size

    def get_count_edges(self):
        return self.gr.number_of_edges()

    def get_in_degrees(self, ind, un=True, w=None):
        try:
            res = [x[0] for x in self.gr.in_edges(ind) if self.get_edge_data(x[0], x[1], w) > 0]
        except:
            res = [x[0] for x in self.gr.in_edges(ind)]

        if un:
            return np.unique(res)
        else:
            return np.array(res)

    def get_out_degrees(self, out, un=True, counts=False, w=None):
        try:
            res = [x[1] for x in self.gr.out_edges(out) if self.get_edge_data(x[0], x[1], w) > 0]
        except:
            res = [x[1] for x in self.gr.out_edges(out)]

        if un:
            return np.unique(res, return_counts=counts)
        else:
            return np.array(res)

    def get_degrees(self, ind, un=True, counts=False, w=None):
        try:
            res = [x[0] for x in self.gr.edges(ind) if
                   self.get_edge_data(x[0], x[1], w) > 0 and x[0] != ind] + [x[1] for x in self.gr.edges(ind) if
                                                                             self.get_edge_data(x[0], x[1], w) > 0 and
                                                                             x[1] != ind]
        except:
            res = [x[0] for x in self.gr.edges(ind) if x[0] != ind] + [x[1] for x in self.gr.edges(ind) if
                                                                       x[1] != ind]

        if un:
            return np.unique(res, return_counts=counts)
        else:
            return np.array(res)

    def update(self, frm, to, node):
        self.gr.node[node][to] = self.gr.node[node][frm]

    def get_data(self, attrs=None, nodes=None):
        if nodes is None:
            nodes = self.gr.nodes()

        if attrs is None:
            attrs = (self.gr.node[nodes[0]]).keys()

        data = []
        for attr in attrs:
            data.append((attr, []))
        ids = []

        for node in nodes:
            ids.append(node)
            for data_line in data:
                data_line[1].append(self.gr.node[node][data_line[0]])

        df = pd.DataFrame(ids, columns=["id"])
        for data_line in data:
            df[data_line[0]] = pd.Series(data_line[1], index=df.index)

        return df

    def save_attrs(self, path, sep=","):
        data = self.get_data()
        data.to_csv(path, sep=sep, index=False)

    def load_attrs(self, path, st=False):
        data = pd.read_csv(path)
        cols = list(data.columns.values)

        for col in cols:
            if col == "id":
                continue
            print(col)
            if st:
                dt = {str(data["id"][i]): str(data[col][i]) for i in np.arange(len(data[col]))}
            else:
                dt = {str(data["id"][i]): convert(data[col][i]) for i in np.arange(len(data[col]))}
            nx.set_node_attributes(self.gr, values=dt, name=col)

    def estimated(self, ind, frm):
        return self.get_attr(ind, frm) >= 0

    def del_attrs(self, attrs):
        if not hasattr(attrs, '__iter__'):
            attrs = [attrs]

        for n in self.get_ids():
            for attr in attrs:
                del self.gr.node[n][attr]

    def get_blk_attrs(self, frm, nodes, typ, part=True):
        xs = np.array([])
        blks = np.array([])

        count = 0
        if typ == "sl" or typ == "sl-ext":
            for j in nodes:
                if self.directed():
                    subs = self.get_in_degrees(j)
                else:
                    subs = self.get_degrees(j)
                if typ == "sl-ext":
                    subs = np.append(j, subs)
                if part:
                    subs = subs[np.isin(subs, nodes)]
                if subs.size == 0:
                    continue
                res = self.get_attributes(frm, subs)
                blks = np.append(blks, count * np.ones(res.size))
                xs = np.append(xs, res)
                count = count + 1

        if typ == "b-sl" or typ == "b-sl-ext":
            for j in nodes:
                subs = self.get_out_degrees(j)
                if typ == "b-sl-ext":
                    subs = np.append(j, subs)
                if part:
                    subs = subs[np.isin(subs, nodes)]
                if subs.size == 0:
                    continue
                res = self.get_attributes(frm, subs)
                blks = np.append(blks, count * np.ones(res.size))
                xs = np.append(xs, res)
                count = count + 1

        if typ == "dj":
            subs = []
            blk = []
            subs_ = np.array([])
            blk_ = []

            for j in nodes:
                sub = self.get_out_degrees(j)

                if sub.size > 1:
                    nad = True
                    for prt in np.arange(len(subs)):
                        if np.sum(np.isin(sub, subs[prt])) > 1:
                            nad = False
                            subs[prt] = np.unique(np.append(subs[prt], sub))
                            blks = np.append(blks, blk[prt])
                            break
                    if nad:
                        subs.append(sub)
                        blk.append(count)
                        blks = np.append(blks, count)
                        count += 1
                else:
                    nex = True
                    for prt in np.arange(subs_.size):
                        if subs_[prt] == sub[0]:
                            nex = False
                            blks = np.append(blks, blk_[prt])
                            break

                    if nex:
                        blks = np.append(blks, count)
                        subs_ = np.append(subs_, sub)
                        blk_.append(count)
                        count += 1
                xs = np.append(xs, self.get_attr(j, frm))

        if typ == "dbl":
            for node in nodes:
                rank = self.get_attr(node, frm)
                for k in self.get_in_degrees(node):
                    if k in nodes:
                        blks = np.append(blks, [count, count])
                        xs = np.append(xs, [rank, self.get_attr(k, frm)])
                        count += 1

        if typ == "ln":
            prep = []
            for node in nodes:
                for k in self.get_in_degrees(node):
                    if k in nodes:
                        prep.append([node, k])

            for node in nodes:
                if self.count_out_degree(node, inn=nodes) == self.count_in_degree(node, inn=nodes) == 1:
                    st = end = -1
                    for j in np.arange(len(prep)):
                        if prep[j][0] == node:
                            st = j
                        if prep[j][-1] == node:
                            end = j
                    st_ = prep[st]
                    end_ = prep[end][1:]
                    prep.append(st_ + end_)
                    del prep[max([st, end])]
                    del prep[min([st, end])]

            for group in prep:
                blks = np.append(blks, [count] * len(group))
                xs = np.append(xs, [self.get_attr(node, frm) for node in group])
                count += 1

        if typ == "mix":
            prep1 = []
            for node in nodes:
                for k in self.get_in_degrees(node):
                    if k in nodes:
                        prep1.append([node, k])

            for node in nodes:
                if self.count_out_degree(node) == self.count_in_degree(node) == 1:
                    st = end = -1
                    for j in np.arange(len(prep1)):
                        if prep1[j][0] == node:
                            st = j
                        if prep1[j][1] == node:
                            end = j
                    st_ = prep1[st]
                    end_ = prep1[end][1:]
                    prep1.append(st_ + end_)
                    del prep1[max([st, end])]
                    del prep1[min([st, end])]

            prep2 = [x for x in nx.cycle_basis(self.subgraph(nodes).gr.to_undirected()) if len(x) == 3]
            prep_ = []
            for group1 in prep1:
                ld = True
                for group2 in prep2:
                    if set(group1).issubset(set(group2)):
                        ld = False
                        break

                if ld:
                    prep_.append(group1)

            prep1 = prep_
            prep = prep1 + prep2
            for group in prep:
                blks = np.append(blks, [count] * len(group))
                xs = np.append(xs, [self.get_attr(node, frm) for node in group])
                count += 1

        return xs, blks

    def get_t_matrix(self, way):
        res = self.get_matrix()
        if not way:
            res = res.transpose()
        return res / (np.sum(res, axis=0) + np.spacing(0))

    def get_matrix(self):
        return nx.to_numpy_matrix(self.get_nx_graph())

    def get_nodes_with_attr(self, attr, value=None, greater=None, less=None):
        if value is None and greater is None and less is None:
            return self.get_ids()

        if not hasattr(value, "__iter__"):
            value = [value]

        if not hasattr(greater, "__iter__"):
            greater = [greater]

        if not hasattr(less, "__iter__"):
            less = [less]

        res = []
        for i in self.get_ids():
            val = self.get_attr(i, attr)

            if val in value:
                res.append(i)
            for g in greater:
                if g is None:
                    break
                if val > g:
                    res.append(i)
                    break
            for l in less:
                if l is None:
                    break
                if val < l:
                    res.append(i)
                    break

        return res


def convert(value):
    try:
        return float(value)
    except:
        return str(value)
