# pylint: disable=missing-docstring

import logging

import nengo
from nengo.builder.operator import Reset
import numpy as np
import pytest
import tensorflow as tf

from nengo_dl import tensor_graph, utils, graph_optimizer, config
from nengo_dl.tests import dummies


@pytest.mark.parametrize("unroll", (1, 2))
@pytest.mark.training
def test_gradients(Simulator, unroll, seed):
    minibatch_size = 4

    with nengo.Network(seed=seed) as net:
        net.config[nengo.Ensemble].gain = nengo.dists.Choice([1])
        net.config[nengo.Ensemble].bias = nengo.dists.Uniform(-1, 1)

        inp = nengo.Node([0], label="inp")

        # sigmoid neurons
        ens = nengo.Ensemble(10, 1, neuron_type=nengo.Sigmoid())

        # normal decoded connection
        nengo.Connection(inp, ens)

        # recurrent connection
        nengo.Connection(ens, ens, transform=0.1)

        # rectified neurons
        ens2 = nengo.Ensemble(10, 2, neuron_type=nengo.RectifiedLinear())

        # neuron--neuron connection
        nengo.Connection(
            ens, ens2, transform=[[1], [1]], solver=nengo.solvers.LstsqL2(weights=True)
        )

        # sliced output, no synapse
        nengo.Connection(inp, ens2[0], synapse=None, transform=0.5)

        # sliced input, sliced output
        inp2 = nengo.Node([0, 0], label="inp2")
        nengo.Connection(inp2[0], ens2[1])

        nengo.Probe(ens)
        nengo.Probe(ens2)

    with Simulator(net, unroll_simulation=unroll, minibatch_size=minibatch_size) as sim:
        sim.check_gradients(atol=1e-4)


def test_mark_signals():
    with nengo.Network() as net:
        ens0 = nengo.Ensemble(10, 1, neuron_type=nengo.LIF())
        ens1 = nengo.Ensemble(20, 1, neuron_type=nengo.Direct())
        ens2 = nengo.Ensemble(30, 1)
        conn0 = nengo.Connection(ens0, ens1)
        conn1 = nengo.Connection(ens0, ens1, learning_rule_type=nengo.PES())
        conn2 = nengo.Connection(ens0, ens2, learning_rule_type=nengo.Voja())
        nengo.Probe(ens2)

    model = nengo.builder.Model()
    model.build(net)

    tg = tensor_graph.TensorGraph(
        model, None, None, 1, None, utils.NullProgressBar(), None
    )
    tg.mark_signals()

    assert model.sig[ens0]["encoders"].trainable
    assert model.sig[ens1]["encoders"].trainable
    assert not model.sig[ens2]["encoders"].trainable
    assert model.sig[ens0.neurons]["bias"].trainable
    assert model.sig[ens2.neurons]["bias"].trainable
    assert model.sig[conn0]["weights"].trainable
    assert not model.sig[conn1]["weights"].trainable
    assert model.sig[conn2]["weights"].trainable

    trainables = (
        model.sig[ens0]["encoders"],
        model.sig[ens1]["encoders"],
        model.sig[ens0.neurons]["bias"],
        model.sig[ens2.neurons]["bias"],
        model.sig[conn0]["weights"],
        model.sig[conn2]["weights"],
    )

    for op in model.operators:
        for sig in op.all_signals:
            if sig in trainables:
                assert sig.trainable
            else:
                assert not sig.trainable


def test_mark_signals_config():
    with nengo.Network() as net:
        config.configure_settings(trainable=None)
        net.config[nengo.Ensemble].trainable = False

        with nengo.Network():
            # check that object in subnetwork inherits config from parent
            ens0 = nengo.Ensemble(10, 1, label="ens0")

            # check that ens.neurons can be set independent of ens
            net.config[ens0.neurons].trainable = True

            with nengo.Network():
                with nengo.Network():
                    # check that subnetworks can override parent configs
                    config.configure_settings(trainable=True)
                    ens1 = nengo.Ensemble(10, 1, label="ens1")

                    with nengo.Network():
                        # check that subnetworks inherit the trainable settings
                        # from parent networks
                        ens3 = nengo.Ensemble(10, 1, label="ens3")

            # check that instances can be set independent of class
            ens2 = nengo.Ensemble(10, 1, label="ens2")
            net.config[ens2].trainable = True

    model = nengo.builder.Model()
    model.build(net)

    progress = utils.NullProgressBar()

    tg = tensor_graph.TensorGraph(model, None, None, 1, None, progress, None)
    tg.mark_signals()

    assert not model.sig[ens0]["encoders"].trainable
    assert model.sig[ens0.neurons]["bias"].trainable

    assert model.sig[ens1]["encoders"].trainable

    assert model.sig[ens2]["encoders"].trainable

    assert model.sig[ens3]["encoders"].trainable

    # check that learning rule connections can be manually set to True
    with nengo.Network() as net:
        config.configure_settings(trainable=None)

        a = nengo.Ensemble(10, 1)
        b = nengo.Ensemble(10, 1)
        conn0 = nengo.Connection(a, b, learning_rule_type=nengo.PES())
        net.config[conn0].trainable = True

    model = nengo.builder.Model()
    model.build(net)

    tg = tensor_graph.TensorGraph(model, None, None, 1, None, progress, None)
    with pytest.warns(UserWarning):
        tg.mark_signals()

    assert model.sig[conn0]["weights"].trainable

    with nengo.Network() as net:
        config.configure_settings(trainable=None)

        a = nengo.Node([0])
        ens = nengo.Ensemble(10, 1)
        nengo.Connection(a, ens, learning_rule_type=nengo.Voja())
        net.config[nengo.Ensemble].trainable = True

    model = nengo.builder.Model()
    model.build(net)

    tg = tensor_graph.TensorGraph(model, None, None, 1, None, progress, None)
    with pytest.warns(UserWarning):
        tg.mark_signals()

    assert model.sig[ens]["encoders"].trainable

    # check that models with no toplevel work
    sig = nengo.builder.signal.Signal([0])
    op = nengo.builder.operator.Reset(sig, 1)
    model = nengo.builder.Model()
    model.add_op(op)

    tg = tensor_graph.TensorGraph(model, None, None, 1, None, progress, None)
    with pytest.warns(UserWarning):
        tg.mark_signals()

    assert not sig.trainable


@pytest.mark.parametrize("config_planner", (True, False, None))
def test_planner_config(config_planner):
    with nengo.Network() as net:
        if config_planner is not None:
            net.config.configures(nengo.Network)
            if config_planner:
                net.config[nengo.Network].set_param(
                    "planner",
                    nengo.params.Parameter("planner", graph_optimizer.noop_planner),
                )

    model = nengo.builder.Model()
    model.build(net)
    sig = nengo.builder.signal.Signal([1])
    sig2 = nengo.builder.signal.Signal([1])
    sig3 = nengo.builder.signal.Signal([1])
    model.add_op(nengo.builder.operator.DotInc(sig, sig2, sig3))
    model.add_op(nengo.builder.operator.DotInc(sig, sig2, sig3))

    tg = tensor_graph.TensorGraph(
        model, None, None, 1, None, utils.NullProgressBar(), None
    )

    assert len(tg.plan) == (3 if config_planner else 2)


def test_signal_order_deterministic(Simulator, seed):
    with nengo.Network(seed=seed) as net:
        inp = nengo.Node([0])
        ens0 = nengo.Ensemble(10, 1, label="ens")
        ens1 = nengo.Ensemble(10, 1, label="ens")
        nengo.Connection(inp, ens0, synapse=None)
        nengo.Connection(inp, ens1, synapse=None)

    with Simulator(net, seed=seed) as sim1:
        with Simulator(net, seed=seed) as sim2:
            for trainable in ("trainable", "non_trainable", "state"):
                for v, v2 in zip(
                    sim1.tensor_graph.base_arrays_init[trainable].values(),
                    sim2.tensor_graph.base_arrays_init[trainable].values(),
                ):
                    assert all(
                        (x is None and y is None) or np.allclose(x, y)
                        for x, y in zip(v[0], v2[0])
                    )


def test_create_signals():
    # check that floats/ints get split into different arrays

    sigs = [
        dummies.Signal(dtype=np.float32),
        dummies.Signal(dtype=np.float32),
        dummies.Signal(dtype=np.int32),
        dummies.Signal(dtype=np.int32),
    ]
    plan = [tuple(dummies.Op(reads=[x]) for x in sigs)]
    graph = dummies.TensorGraph(plan, tf.float32, 10)
    graph.create_signals(sigs)

    assert graph.signals[sigs[0]].key == graph.signals[sigs[1]].key
    assert graph.signals[sigs[1]].key != graph.signals[sigs[2]].key
    assert graph.signals[sigs[2]].key == graph.signals[sigs[3]].key

    # check that floats all get converted to same precision and combined
    sigs = [
        dummies.Signal(dtype=np.float32),
        dummies.Signal(dtype=np.float32),
        dummies.Signal(dtype=np.float64),
        dummies.Signal(dtype=np.float64),
    ]
    plan = [tuple(dummies.Op(reads=[x]) for x in sigs)]
    graph = dummies.TensorGraph(plan, tf.float32, 10)
    graph.create_signals(sigs)
    assert np.all([graph.signals[x].dtype == "float32" for x in sigs])
    assert graph.signals[sigs[0]].key == graph.signals[sigs[1]].key
    assert graph.signals[sigs[1]].key == graph.signals[sigs[2]].key
    assert graph.signals[sigs[2]].key == graph.signals[sigs[3]].key

    # check that ints all get converted to same precision and combined
    sigs = [
        dummies.Signal(dtype=np.int32),
        dummies.Signal(dtype=np.int32),
        dummies.Signal(dtype=np.int64),
        dummies.Signal(dtype=np.int64),
    ]
    plan = [tuple(dummies.Op(reads=[x]) for x in sigs)]
    graph = dummies.TensorGraph(plan, tf.float32, 10)
    graph.create_signals(sigs)
    assert np.all([graph.signals[x].dtype == "int32" for x in sigs])
    assert graph.signals[sigs[0]].key == graph.signals[sigs[1]].key
    assert graph.signals[sigs[1]].key == graph.signals[sigs[2]].key
    assert graph.signals[sigs[2]].key == graph.signals[sigs[3]].key

    # check that different shapes go in different groups
    sigs = [
        dummies.Signal(shape=(10,)),
        dummies.Signal(shape=(5,)),
        dummies.Signal(shape=(10, 1)),
        dummies.Signal(shape=(5, 1)),
    ]
    plan = [tuple(dummies.Op(reads=[x]) for x in sigs)]
    graph = dummies.TensorGraph(plan, tf.float32, 10)
    graph.create_signals(sigs)
    assert graph.base_arrays_init["non_trainable"][graph.signals[sigs[0]].key][1] == [
        (10, 10),
        (10, 5),
    ]
    assert graph.base_arrays_init["non_trainable"][graph.signals[sigs[2]].key][1] == [
        (10, 10, 1),
        (10, 5, 1),
    ]
    assert graph.signals[sigs[0]].key == graph.signals[sigs[1]].key
    assert graph.signals[sigs[1]].key != graph.signals[sigs[2]].key
    assert graph.signals[sigs[2]].key == graph.signals[sigs[3]].key

    # check trainable
    sigs = [
        dummies.Signal(trainable=True),
        dummies.Signal(trainable=True),
        dummies.Signal(trainable=False),
        dummies.Signal(trainable=False),
    ]
    plan = [tuple(dummies.Op(reads=[x]) for x in sigs)]
    graph = dummies.TensorGraph(plan, tf.float32, 10)
    graph.create_signals(sigs)
    assert graph.base_arrays_init["trainable"][graph.signals[sigs[0]].key][1] == [
        (1,),
        (1,),
    ]
    assert graph.base_arrays_init["non_trainable"][graph.signals[sigs[2]].key][1] == [
        (10, 1),
        (10, 1),
    ]
    assert graph.signals[sigs[0]].key == graph.signals[sigs[1]].key
    assert graph.signals[sigs[1]].key != graph.signals[sigs[2]].key
    assert graph.signals[sigs[2]].key == graph.signals[sigs[3]].key

    # check that scalars get upsized
    sigs = [dummies.Signal(shape=()), dummies.Signal(shape=(4,))]
    plan = [tuple(dummies.Op(reads=[x]) for x in sigs)]
    graph = dummies.TensorGraph(plan, tf.float32, 10)
    graph.create_signals(sigs)
    assert list(graph.base_arrays_init["non_trainable"].values())[0][1] == [
        (10, 1),
        (10, 4),
    ]

    # check that boolean signals are handled correctly
    sigs = [dummies.Signal(dtype=np.bool, shape=())]
    plan = [(dummies.Op(reads=sigs),)]
    graph = dummies.TensorGraph(plan, tf.float32, 1)
    graph.create_signals(sigs)
    assert list(graph.base_arrays_init["non_trainable"].values())[0][2] == "bool"


def test_create_signals_views():
    sigs = [
        dummies.Signal(shape=(2, 2), base_shape=(4,)),
        dummies.Signal(shape=(2, 2), base_shape=(4,)),
    ]
    sigs += [sigs[0].base, sigs[1].base]
    plan = [tuple(dummies.Op(reads=[x]) for x in sigs)]
    graph = dummies.TensorGraph(plan, tf.float32, 10)
    graph.create_signals(sigs[2:])
    assert list(graph.base_arrays_init["non_trainable"].values())[0][1] == [
        (10, 4),
        (10, 4),
    ]
    assert graph.signals[sigs[0]].key == graph.signals[sigs[1]].key
    assert graph.signals[sigs[1]].key == graph.signals[sigs[2]].key
    assert graph.signals[sigs[2]].key == graph.signals[sigs[3]].key
    assert graph.signals[sigs[0]].slices == ((0, 4),)
    assert graph.signals[sigs[1]].slices == ((4, 8),)
    assert graph.signals[sigs[0]].slices == graph.signals[sigs[2]].slices
    assert graph.signals[sigs[1]].slices == graph.signals[sigs[3]].slices


def test_create_signals_partition():
    # check that signals are partitioned based on plan
    sigs = [dummies.Signal(), dummies.Signal(), dummies.Signal(), dummies.Signal()]
    plan = [
        tuple(dummies.Op(reads=[x]) for x in sigs[:2]),
        tuple(dummies.Op(reads=[x]) for x in sigs[2:]),
    ]
    graph = dummies.TensorGraph(plan, tf.float32, 10)
    graph.create_signals(sigs)
    assert graph.signals[sigs[0]].key == graph.signals[sigs[1]].key
    assert graph.signals[sigs[1]].key != graph.signals[sigs[2]].key
    assert graph.signals[sigs[2]].key == graph.signals[sigs[3]].key

    # check that signals are partitioned for different read blocks
    plan = [tuple(dummies.Op(reads=[sigs[i], sigs[2 + i]]) for i in range(2))]
    graph = dummies.TensorGraph(plan, tf.float32, 10)
    graph.create_signals(sigs)
    assert graph.signals[sigs[0]].key == graph.signals[sigs[1]].key
    assert graph.signals[sigs[1]].key != graph.signals[sigs[2]].key
    assert graph.signals[sigs[2]].key == graph.signals[sigs[3]].key

    # check that signals are partitioned for different sig types
    plan = [tuple(dummies.Op(reads=[sigs[i]], sets=[sigs[2 + i]]) for i in range(2))]
    graph = dummies.TensorGraph(plan, tf.float32, 10)
    graph.create_signals(sigs)
    assert graph.signals[sigs[0]].key == graph.signals[sigs[1]].key
    assert graph.signals[sigs[1]].key != graph.signals[sigs[2]].key
    assert graph.signals[sigs[2]].key == graph.signals[sigs[3]].key

    # check that resets are ignored
    sigs = [dummies.Signal(), dummies.Signal(), dummies.Signal(), dummies.Signal()]
    plan = [tuple(Reset(x) for x in sigs)]
    graph = dummies.TensorGraph(plan, tf.float32, 10)
    graph.create_signals(sigs)
    assert len(graph.base_arrays_init["non_trainable"]) == 4


@pytest.mark.parametrize("use_loop", (True, False))
def test_get_tensor(Simulator, use_loop):
    with nengo.Network() as net:
        config.configure_settings(use_loop=use_loop)

        a = nengo.Node([1])
        b = nengo.Ensemble(10, 1)
        c = nengo.Connection(
            a, b.neurons, transform=np.arange(10)[:, None], synapse=None
        )
        p = nengo.Probe(c)

        # build a signal probe so that the indices get loaded into the sim
        # (checks that the indices reloading works properly)
        nengo.Probe(c, "weights")

    kwargs = dict() if use_loop else dict(unroll_simulation=10)
    with Simulator(net, **kwargs) as sim:
        tensor = sim.tensor_graph.get_tensor(sim.model.sig[c]["weights"])

        assert np.allclose(tf.keras.backend.get_value(tensor), np.arange(10)[:, None])

        sim.run_steps(10)
        assert np.allclose(sim.data[p], np.arange(10)[None, :])


@pytest.mark.parametrize("trainable", (True, False))
def test_build(trainable, rng):
    sigs = [
        dummies.Signal(
            shape=(2, 1), dtype="float32", initial_value=0, trainable=trainable
        ),
        dummies.Signal(
            shape=(3, 1),
            dtype="float32",
            initial_value=np.zeros((3, 1)),
            trainable=trainable,
        ),
        dummies.Signal(
            shape=(4, 1), dtype="float32", initial_value=1, trainable=trainable
        ),
        dummies.Signal(
            shape=(5, 1),
            dtype="float32",
            initial_value=np.ones((5, 1)),
            trainable=trainable,
        ),
        dummies.Signal(
            shape=(6, 1),
            dtype="float32",
            initial_value=rng.uniform(size=(6, 1)),
            trainable=trainable,
        ),
        dummies.Signal(
            shape=(7, 1),
            dtype="float32",
            initial_value=rng.uniform(size=(7, 1)),
            trainable=trainable,
        ),
    ]

    plan = [
        tuple(dummies.Op(reads=[x]) for x in sigs[:2]),
        tuple(dummies.Op(reads=[x]) for x in sigs[2:4]),
        tuple(dummies.Op(reads=[x]) for x in sigs[4:]),
    ]

    graph = dummies.TensorGraph(plan=plan, dtype="float32", minibatch_size=16)
    graph.create_signals(sigs)
    graph.build()

    if trainable:
        assert len(graph.trainable_weights) == 3
        assert len(graph.non_trainable_weights) == 0
    else:
        assert len(graph.trainable_weights) == 0
        assert len(graph.non_trainable_weights) == 3

    init0 = graph.weights[0].numpy()
    assert init0.shape == (5, 1) if trainable else (16, 5, 1)
    assert np.allclose(init0, 0)

    init1 = graph.weights[1].numpy()
    assert init1.shape == (9, 1) if trainable else (16, 9, 1)
    assert np.allclose(init1, 1)

    init2 = graph.weights[2].numpy()
    if trainable:
        assert init2.shape == (13, 1)
        assert np.allclose(init2[:6], sigs[4].initial_value)
        assert np.allclose(init2[6:], sigs[5].initial_value)
    else:
        assert init2.shape == (16, 13, 1)
        assert np.allclose(init2[:, :6], sigs[4].initial_value)
        assert np.allclose(init2[:, 6:], sigs[5].initial_value)


@pytest.mark.parametrize("use_loop", (True, False))
def test_conditional_update(Simulator, use_loop, caplog):
    caplog.set_level(logging.INFO)

    with nengo.Network() as net:
        config.configure_settings(stateful=False, use_loop=use_loop)

        a = nengo.Ensemble(10, 1)
        b = nengo.Node(size_in=1)
        conn = nengo.Connection(a, b)

    with Simulator(net):
        pass

    assert "Number of variable updates: 0" in caplog.text
    caplog.clear()

    conn.learning_rule_type = nengo.PES()

    with Simulator(net):
        pass

    assert "Number of variable updates: 1" in caplog.text
    caplog.clear()

    with net:
        config.configure_settings(trainable=True)

    with Simulator(net):
        pass

    assert "Number of variable updates: 1" in caplog.text
