import nengo
import numpy as np
import pytest

import nengo_spa as spa
from nengo_spa.algebras.vtb_algebra import VtbAlgebra
from nengo_spa.networks.vtb import VTB
from nengo_spa.testing import assert_sp_close


def test_bind(Simulator, seed):
    rng = np.random.RandomState(seed)
    vocab = spa.Vocabulary(16, pointer_gen=rng, algebra=VtbAlgebra())
    vocab.populate('A; B')

    with spa.Network(seed=seed) as model:
        vtb = VTB(100, 16)
        nengo.Connection(nengo.Node(vocab['A'].v), vtb.input_left)
        nengo.Connection(nengo.Node(vocab['B'].v), vtb.input_right)
        p = nengo.Probe(vtb.output, synapse=0.03)

    with Simulator(model) as sim:
        sim.run(0.2)

    assert_sp_close(
        sim.trange(), sim.data[p], vocab.parse('A*B'), skip=0.15, atol=0.3)


@pytest.mark.parametrize('side', ('left', 'right'))
def test_unbind(Simulator, side, seed):
    rng = np.random.RandomState(seed)
    vocab = spa.Vocabulary(36, pointer_gen=rng, algebra=VtbAlgebra())
    vocab.populate('A; B')

    with spa.Network(seed=seed) as model:
        vtb = VTB(
            100, 36, unbind_left=(side == 'left'),
            unbind_right=(side == 'right'))

        if side == 'left':
            left = nengo.Node(vocab['B'].v)
            right = nengo.Node(vocab.parse('B*A').v)
        elif side == 'right':
            left = nengo.Node(vocab.parse('A*B').v)
            right = nengo.Node(vocab['B'].v)
        else:
            raise ValueError("Invalid 'side' value.")

        nengo.Connection(left, vtb.input_left)
        nengo.Connection(right, vtb.input_right)

        p = nengo.Probe(vtb.output, synapse=0.03)

    with Simulator(model) as sim:
        sim.run(0.2)

    assert_sp_close(
        sim.trange(), sim.data[p], vocab.parse('A * B * ~B'), skip=0.15,
        atol=0.3)
