from brian import *
from nose.tools import *

def test():
    '''
    Equations module.
    '''
    reinit_default_clock()

    tau = 2 * ms

    # Test Parsing and building
    def test_build_and_parse(eqs):
        assert set(eqs._eq_names) == set(['x', 'y']) # An alias is also a static equation
        assert set(eqs._diffeq_names) == set(['v', 'z']) # A parameter is a differential equation dz/dt=0
        assert eqs._diffeq_names_nonzero == ['v']
        assert eqs._alias == {'y':'x'}
        assert eqs._units == {'t':second, 'x':volt ** 2, 'z':amp, 'y':volt ** 2, 'v':volt}
        assert eqs._string == {'y':'x', 'x':'v**2', 'z':'0*amp/second', 'v':'(y-v)/tau'}
        assert eqs._namespace['v']['tau'] == 2 * ms
        assert 'tau' not in eqs._namespace['x']

    # Multiline string
    eqs = Equations('''
    x=v**2 : volt**2
    y=x # a comment
    dv/dt=(y-v)/tau : volt
    z : amp''')
    test_build_and_parse(eqs)
    
    # Concatenated Equation objects
    eqs = (Equations('''x=v**2 : volt**2
                        y=x # a comment''') +
           Equations('dv/dt=(y-v)/tau : volt') +
           Equations('z : amp'))
    test_build_and_parse(eqs)
    
    # List of Equation objects
    eqs = Equations([Equations('''x=v**2 : volt**2
                                  y=x # a comment'''),
                     Equations('dv/dt=(y-v)/tau : volt'),
                     Equations('z : amp')])
    test_build_and_parse(eqs)
    
    # List of strings
    eqs = Equations(['''x=v**2 : volt**2
                        y=x # a comment''',
                     'dv/dt=(y-v)/tau : volt',
                     'z : amp'])
    test_build_and_parse(eqs)
    
    # Name substitutions
    assert Equations('dx/dt=-x/(2*ms) : 1', x='y')._diffeq_names == ['y']
    assert Equations('''dx/dt=-x/(2*ms) : 1
                        z = x ** 2 : 1''', z='y')._eq_names == ['y']

    # Explicit namespace
    eqs2 = Equations('dx/dt=-x/tau:volt', tau=1 * ms)
    assert eqs2._namespace['x']['tau'] == 1 * ms
    assert eqs2._namespace['x'] == {'tau':1 * ms, 'volt':volt}

    # Find membrane potential
    assert eqs.get_Vm() == 'v'
    assert Equations('v=x**2 : 1').get_Vm() == None # must be a differential equation
    assert Equations('dx/dt=1/(2*ms) : 1').get_Vm() == None
    assert Equations('dvm/dt=1/(2*ms) : 1').get_Vm() == 'vm'
    assert Equations('''dvm/dt=1/(2*ms) : 1
    dv/dt=1/(2*ms) : 1
    ''').get_Vm() == None # ambiguous

    # Unit checking
    eqs = Equations('dv/dt=-v : volt')
    assert_raises(DimensionMismatchError, eqs.prepare)

    # Free variables
    eqs = Equations('dv/dt=(freevar-v)/tau : volt')
    assert eqs.free_variables() == ['tau', 'freevar']

    # Equation types: stochastic, time dependent, linear, conditionally linear
    eqs = Equations('''
    dv/dt=(f(t)-v)/tau : volt
    dw/dt=-w/tau+3*tau**-.5*xi : 1''')
    assert eqs.is_stochastic()
    assert eqs.is_stochastic(var='v') is False
    assert eqs.is_stochastic(var='w')
    assert eqs.is_time_dependent()
    assert eqs.is_time_dependent(var='v')
    assert eqs.is_time_dependent(var='w') is False
    assert eqs.is_linear() is False # not linear if time-dependent

    eqs = Equations('dv/dt=-v/tau : volt')
    assert eqs.is_linear()
    eqs.prepare()
    assert eqs.is_conditionally_linear()

    eqs = Equations('''
    dv/dt=(w**3-v)/tau : 1
    dw/dt=(v**2-w)/tau : 1
    ''')
    eqs.prepare()
    assert eqs.is_conditionally_linear()

    eqs = Equations('''
    dv/dt=(w**3-v**2)/tau : 1
    dw/dt=(v**2-w)/tau : 1
    ''')
    eqs.prepare()
    assert eqs.is_conditionally_linear() is False

if __name__ == '__main__':
    test()
