import networkx as nx
import numpy as np
import jax
import jax.numpy as jnp
from functools import partial
from scipy.optimize import minimize
from jax.experimental.sparse import BCOO
from spinn import utils, graphs

# import os
# os.environ['CUDA_VISIBLE_DEVICES'] = ''
# jax.config.update("jax_enable_x64", True)
# jax.config.update('jax_default_matmul_precision', 'high')  # 'bfloat16_3x'


class spinn:
    def __init__(self, G, f=None, p=None,
                 initx=None, inity=None,
                 key=None,
                 sparse=False):
        if G is None:
            return  # for tree unflattenning
        if isinstance(G, (list, tuple, np.ndarray)):
            G = graphs.mlgraph(G)
        assert nx.is_directed_acyclic_graph(G)
        layers = [sorted(g) for g in nx.topological_generations(G)]
        nlst = sum(layers, [])
        H = nx.relabel_nodes(G, dict(zip(nlst, range(len(G)))))
        layers = [sorted(g) for g in nx.topological_generations(H)]
        arch = [len(li) for li in layers]
        carch = utils.acc(arch)  # list(itertools.accumulate(arch))
        sizes = list(zip(arch[1:], carch[:-1]))  # sizes of weight matrices
        conecs = [list(zip(*H.in_edges(li))) for li in layers[1:]]
        wmrk = utils.pacc([0] + [len(ci[0]) for ci in conecs])
        bmrk = utils.pacc([wmrk[-1][1]] + arch[1:])
        inps = sorted((n for n, d in G.in_degree() if d == 0))
        outs = sorted((n for n, d in G.out_degree() if d == 0))
        f, p = utils.unroll_fp(f, p, arch)
        self.nlst = nlst
        self.n = len(H)
        self.inps = inps
        self.outs = outs
        self.ni = len(self.inps)
        self.no = len(self.outs)
        self.nw = len(H.edges)
        self.nb = self.n - self.ni
        self.wshp = [(s, (tuple(np.asarray(i1)-i1[0]), i0), wm, bm)
                     for s, (i0, i1), wm, bm in zip(sizes, conecs, wmrk, bmrk)]
        self.f = f
        self.p = p
        self.sparse = sparse
        self.initw(key=key)
        if initx is not None:
            self.initx(initx)
        if inity is not None:
            self.inity(inity)
        # self.codes = {None: [[i] for i in range(self.ni)],
        #               'F': [[0]], 'dFx': [[0, 0]]}
        self._register_pytree()

    def initw(self, key=None):
        w = []
        b = []
        for s, (i0, i1), _, _ in self.wshp:
            u, c = np.unique(i0, return_counts=True)
            n = np.repeat(c, c)
            state = np.random.get_state()
            np.random.seed(key)
            r = np.random.uniform(-1., 1., size=len(i0)+s[0])
            np.random.set_state(state)
            w += [jnp.asarray(2.38*r[:len(i0)]/np.sqrt(n+1))]
            b += [jnp.asarray(2.38*r[len(i0):]/np.sqrt(c+1))]
        w = jnp.concatenate(w + b)
        self.w = w
        return w

    def initx(self, x):
        x = jnp.atleast_2d(x)
        assert x.shape[1] == self.ni
        ax, bx = utils.to01c(x)
        self.f[0] = 'linear'
        self.p[0] = (ax, bx)
        return ax, bx

    def inity(self, y):
        y = jnp.atleast_2d(y)
        assert y.shape[1] == self.no
        ay, by = utils.from01c(y)
        self.f[-1] = 'linear'
        self.p[-1] = (ay, by)
        return ay, by

    def afunc(self, xi, i):
        f = self.f[i]
        p = self.p[i]
        match f:
            case 'identity':
                return xi
            case 'linear':
                return self.linear(xi, *p)
            case 'relu':
                return jax.nn.relu(xi, *p)
            case 'sigmoid':
                return jax.nn.sigmoid(xi, *p)
            case 'softplus':
                return jax.nn.softplus(xi, *p)
            case _:
                return f(xi, *p)

    def _prop(self, x):
        y = self.afunc(x, 0)
        Y = jnp.zeros(self.n).at[jnp.asarray(self.inps)].set(y)
        for i, (s, (i0, i1), (j, k), (l, m)) in enumerate(self.wshp):
            wi = self.w[j:k]
            bi = self.w[l:m]
            if self.sparse:
                wi = BCOO((wi, jnp.asarray((i0, i1)).T), shape=s)
            else:
                wi = jnp.zeros(s).at[i0, i1].set(wi)  # dense is default
            yi = Y[:s[1]]
            xi = wi@yi + bi
            y = self.afunc(xi, i+1)
            Y = Y.at[s[1]:s[1]+s[0]].set(y)
        return Y[jnp.asarray(self.outs)]

    @jax.jit
    @partial(jax.vmap, in_axes=(None, 0, None))
    def propd(self, x, spec):
        f = self._prop

        def deriv(f, x, v):
            return jax.jvp(f, [x], [v])[1]

        def one_hot(i):
            return jnp.zeros(len(x)).at[i].set(1)

        df = []
        dy = []
        for i, a, in enumerate(spec):
            df += [[f]]
            for j, n in enumerate(a[1:]):
                def dfij(x, i=i, j=j, n=n):
                    vn = one_hot(n)
                    return deriv(df[i][j], x, vn)
                df[i] += [dfij]
            k = jnp.atleast_1d(jnp.asarray(a[0]))
            dy += [df[i][-1](x)[k]]
        return jnp.concatenate(dy)

    @jax.jit
    def propx(self, xspec):
        return tuple(self.propd(a[0], [a[1:]]).ravel() for a in xspec)

    @jax.jit
    @partial(jax.vmap, in_axes=(None, 0))
    def __call__(self, x):
        return self._prop(x)

    # @jax.jit
    # def __call__(self, x, spec=None):
    #     spec = self.codes.get(spec, spec)
    #     return self._propd(x, spec)

    @jax.jit
    def loss(self, w, x, t):
        self.w = jnp.asarray(w)
        return self.mse(self(x), t)

    @jax.jit
    @partial(jax.grad, argnums=1)
    def lossg(self, w, *args):
        return self.loss(w, *args)

    @jax.jit
    @partial(jax.value_and_grad, argnums=1)
    def lossvg(self, w, *args):
        return self.loss(w, *args)

    @jax.jit
    @partial(jax.jacobian, argnums=1)
    def eqcj(self, w, *args):
        return self.eqc(w, *args)

    @jax.jit
    @partial(jax.jacobian, argnums=1)
    def ineqcj(self, w, *args):
        return self.ineqc(w, *args)

    # @jax.jit
    # def loss(self, w, *args):
    #     self.w = jnp.asarray(w)
    #     y = self.eqs(w, *args)
    #     return sum(mse(yi) for yi in y)

    # This is concrete represntation of eqs
    # @jax.jit
    # def eqs(self, w, x, t):
    #     self.w = jnp.asarray(w)
    #     return [self(x) - t]

    # Unfortunately jax's BFGS most of the time fails because of 3 - zoom error
    # def train(self, *args, maxiter=None, tol=None, method='BFGS', **kwargs):
    #     from jax.scipy.optimize import minimize
    #     res = minimize(self.loss, self.w, args,
    #                    tol=tol,
    #                    options={'maxiter': maxiter},
    #                    method=method,
    #                    **kwargs)
    #     self.w = jnp.asarray(res.x)
    #     return res

    # We use scipy.optimize.minimize here. It has many advantages,
    # but cannot be jitted...
    def train(self, *args, tol=1e-5, maxiter=None, disp=False, **kwargs):
        options = kwargs.pop('options', {})
        options.setdefault('maxiter', maxiter)
        options.setdefault('disp', disp)
        cstr = []
        if hasattr(self, 'eqc'):
            cstr += [{'type': 'eq',
                      'fun': self.eqc,
                      'jac': self.eqcj,
                      'args': args}]
        if hasattr(self, 'ineqc'):
            cstr += [{'type': 'ineq',
                      'fun': self.ineqc,
                      'jac': self.ineqcj,
                      'args': args}]
        method = 'BFGS' if len(cstr) == 0 else 'SLSQP'
        if method == 'SLSQP':
            options['ftol'] = tol  # it seems this is set here
        res = minimize(self.loss, self.w, args, jac=self.lossg,
                       method=method,
                       tol=tol,
                       constraints=cstr,
                       options=options,
                       **kwargs)
        self.w = jnp.asarray(res.x)
        return res

    @staticmethod
    def linear(x, a=1., b=0.):
        return a*x + b

    @staticmethod
    def mse(y, t=0.):
        y = jnp.asarray(y).ravel()
        t = jnp.asarray(t).ravel()
        return jnp.mean((y - t)**2)

    @staticmethod
    def rmse(y, t=0.):
        return jnp.sqrt(spinn.mse(y, t))

    def _tree_flatten(self):
        children = (self.w, self.p)     # dynamic values
        aux_data = {'nlst': self.nlst,
                    'n': self.n,
                    'inps': self.inps,
                    'outs': self.outs,
                    'ni': self.ni,
                    'no': self.no,
                    'nw': self.nw,
                    'nb': self.nb,
                    'wshp': self.wshp,
                    'f': self.f,
                    'sparse': self.sparse}  # static values
        return (children, aux_data)

    @classmethod
    def _tree_unflatten(cls, aux_data, children):
        obj = cls(None)
        obj.w = children[0]
        obj.p = children[1]
        obj.__dict__.update(aux_data)
        return obj

    @classmethod
    def _register_pytree(cls):
        # This method is called at __init__ and allows for
        # automatic pytree registration of subclasses.
        # Function 'register_pytree_node' returns ValueError when registering
        # the same class many times, so there is the 'if' below. Anyway,
        # in case of troubles try to do this explicitly!
        if cls not in jax._src.tree_util._registry:
            jax.tree_util.register_pytree_node(cls,
                                               cls._tree_flatten,
                                               cls._tree_unflatten)


if __name__ == "__main__":
    import os
    import subprocess
    fname = os.path.dirname(os.path.realpath(__file__))
    subprocess.call(['pytest', fname])
