from collections import namedtuple
import json
import os
import platform
import socket

import numpy as np
import pytest

from brian2 import prefs, clear_cache, _cache_dirs_and_extensions
from brian2.codegen.cpp_prefs import compiler_supports_c99, get_compiler_and_args
from brian2.codegen.optimisation import optimise_statements
from brian2.codegen.runtime.cython_rt import CythonCodeObject
from brian2.codegen.translation import (analyse_identifiers,
                                        get_identifiers_recursively,
                                        parse_statement,
                                        make_statements,
                                        )
from brian2.codegen.statements import Statement
from brian2.codegen.codeobject import CodeObject
from brian2.parsing.sympytools import str_to_sympy, sympy_to_str
from brian2.core.variables import Subexpression, Variable, Constant, ArrayVariable
from brian2.core.functions import Function, DEFAULT_FUNCTIONS, DEFAULT_CONSTANTS
from brian2.devices.device import auto_target, device
from brian2.units.fundamentalunits import Unit
from brian2.units import second, ms
from brian2.utils.logger import catch_logs

FakeGroup = namedtuple('FakeGroup', ['variables'])

@pytest.mark.codegen_independent
def test_auto_target():
    # very basic test that the "auto" codegen target is useable
    assert issubclass(auto_target(), CodeObject)


@pytest.mark.codegen_independent
def test_analyse_identifiers():
    """
    Test that the analyse_identifiers function works on a simple clear example.
    """
    code = """
    a = b+c
    d = e+f
    """
    known = {'b': Variable(name='b'),
             'c': Variable(name='c'),
             'd': Variable(name='d'),
             'g': Variable(name='g')}
    
    defined, used_known, dependent = analyse_identifiers(code, known)
    assert 'a' in defined  # There might be an additional constant added by the
                           # loop-invariant optimisation
    assert used_known == {'b', 'c', 'd'}
    assert dependent == {'e', 'f'}


@pytest.mark.codegen_independent
def test_get_identifiers_recursively():
    """
    Test finding identifiers including subexpressions.
    """
    variables = {'sub1': Subexpression(name='sub1',
                                       dtype=np.float32, expr='sub2 * z',
                                       owner=FakeGroup(variables={}),
                                       device=None),
                 'sub2': Subexpression(name='sub2',
                                       dtype=np.float32, expr='5 + y',
                                       owner=FakeGroup(variables={}),
                                       device=None),
                 'x': Variable(name='x')}
    identifiers = get_identifiers_recursively(['_x = sub1 + x'],
                                              variables)
    assert identifiers == {'x', '_x', 'y', 'z', 'sub1', 'sub2'}


@pytest.mark.codegen_independent
def test_write_to_subexpression():
    variables = {
        'a': Subexpression(name='a', dtype=np.float32,
                           owner=FakeGroup(variables={}), device=None,
                           expr='2*z'),
        'z': Variable(name='z')
    }

    # Writing to a subexpression is not allowed
    code = 'a = z'
    with pytest.raises(SyntaxError):
        make_statements(code, variables, np.float32)


@pytest.mark.codegen_independent
def test_repeated_subexpressions():
    variables = {
        'a': Subexpression(name='a', dtype=np.float32,
                           owner=FakeGroup(variables={}), device=None,
                           expr='2*z'),
        'x': Variable(name='x'),
        'y': Variable(name='y'),
        'z': Variable(name='z')
    }
    # subexpression a (referring to z) is used twice, but can be reused the
    # second time (no change to z)
    code = """
    x = a
    y = a
    """
    scalar_stmts, vector_stmts = make_statements(code, variables, np.float32)
    assert len(scalar_stmts) == 0
    assert [stmt.var for stmt in vector_stmts] == ['a', 'x', 'y']
    assert vector_stmts[0].constant

    code = """
    x = a
    z *= 2
    """
    scalar_stmts, vector_stmts = make_statements(code, variables, np.float32)
    assert len(scalar_stmts) == 0
    assert [stmt.var for stmt in vector_stmts] == ['a', 'x', 'z']
    # Note that we currently do not mark the subexpression as constant in this
    # case, because its use after the "z *=2" line would actually redefine it.
    # Our algorithm is currently not smart enough to detect that it is actually
    # not used afterwards

    # a refers to z, therefore we have to redefine a after z changed, and a
    # cannot be constant
    code = """
    x = a
    z *= 2
    y = a
    """
    scalar_stmts, vector_stmts = make_statements(code, variables, np.float32)
    assert len(scalar_stmts) == 0
    assert [stmt.var for stmt in vector_stmts] == ['a', 'x', 'z', 'a', 'y']
    assert not any(stmt.constant for stmt in vector_stmts)


@pytest.mark.codegen_independent
def test_nested_subexpressions():
    """
    This test checks that code translation works with nested subexpressions.
    """
    code = """
    x = a + b + c
    c = 1
    x = a + b + c
    d = 1
    x = a + b + c
    """
    variables = {
        'a': Subexpression(name='a', dtype=np.float32, owner=FakeGroup(variables={}), device=None,
                           expr='b*b+d'),
        'b': Subexpression(name='b', dtype=np.float32, owner=FakeGroup(variables={}), device=None,
                           expr='c*c*c'),
        'c': Variable(name='c'),
        'd': Variable(name='d'),
        }
    scalar_stmts, vector_stmts = make_statements(code, variables, np.float32)
    assert len(scalar_stmts) == 0
    evalorder = ''.join(stmt.var for stmt in vector_stmts)
    # This is the order that variables ought to be evaluated in (note that
    # previously this test did not expect the last "b" evaluation, because its
    # value did not change (c was not changed). We have since removed this
    # subexpression caching, because it did not seem to apply in practical
    # use cases)
    assert evalorder == 'baxcbaxdbax'

@pytest.mark.codegen_independent
def test_apply_loop_invariant_optimisation():
    variables = {'v': Variable('v', scalar=False),
                 'w': Variable('w', scalar=False),
                 'dt': Constant('dt', dimensions=second.dim, value=0.1*ms),
                 'tau': Constant('tau', dimensions=second.dim, value=10*ms),
                 'exp': DEFAULT_FUNCTIONS['exp']}
    statements = [Statement('v', '=', 'dt*w*exp(-dt/tau)/tau + v*exp(-dt/tau)', '', np.float32),
                  Statement('w', '=', 'w*exp(-dt/tau)', '', np.float32)]
    scalar, vector = optimise_statements([], statements, variables)
    # The optimisation should pull out at least exp(-dt / tau)
    assert len(scalar) >= 1
    assert np.issubdtype(scalar[0].dtype, np.floating)
    assert scalar[0].var == '_lio_1'
    assert len(vector) == 2
    assert all('_lio_' in stmt.expr for stmt in vector)

@pytest.mark.codegen_independent
def test_apply_loop_invariant_optimisation_integer():
    variables = {'v': Variable('v', scalar=False),
                 'N': Constant('N', 10),
                 'b': Variable('b', scalar=True, dtype=int),
                 'c': Variable('c', scalar=True, dtype=int),
                 'd': Variable('d', scalar=True, dtype=int),
                 'y': Variable('y', scalar=True, dtype=float),
                 'z': Variable('z', scalar=True, dtype=float),
                 'w': Variable('w', scalar=True, dtype=float),
                 }
    statements = [Statement('v', '=', 'v % (2*3*N)', '', np.float32),
                  # integer version doesn't get rewritten but float version does
                  Statement('a', ':=', 'b//(c//d)', '', int),
                  Statement('x', ':=', 'y/(z/w)', '', float),
                  ]
    scalar, vector = optimise_statements([], statements, variables)
    assert len(scalar) == 3
    assert np.issubdtype(scalar[0].dtype, np.signedinteger)
    assert scalar[0].var == '_lio_1'
    expr = scalar[0].expr.replace(' ', '')
    assert expr=='6*N' or expr=='N*6'
    assert np.issubdtype(scalar[1].dtype, np.signedinteger)
    assert scalar[1].var == '_lio_2'
    expr = scalar[1].expr.replace(' ', '')
    assert expr=='b//(c//d)'
    assert np.issubdtype(scalar[2].dtype, np.floating)
    assert scalar[2].var == '_lio_3'
    expr = scalar[2].expr.replace(' ', '')
    assert expr=='(y*w)/z' or expr=='(w*y)/z'

@pytest.mark.codegen_independent
def test_apply_loop_invariant_optimisation_boolean():
    variables = {'v1': Variable('v1', scalar=False),
                 'v2': Variable('v2', scalar=False),
                 'N': Constant('N', 10),
                 'b': Variable('b', scalar=True, dtype=bool),
                 'c': Variable('c', scalar=True, dtype=bool),
                 'int': DEFAULT_FUNCTIONS['int'],
                 'foo': Function(lambda x: None,
                                 arg_units=[Unit(1)], return_unit=Unit(1),
                                 arg_types=['boolean'], return_type='float',
                                 stateless=False)
                 }
    # The calls for "foo" cannot be pulled out, since foo is marked as stateful
    statements = [Statement('v1', '=', '1.0*int(b and c)', '', np.float32),
                  Statement('v1', '=', '1.0*foo(b and c)', '', np.float32),
                  Statement('v2', '=', 'int(not b and True)', '', np.float32),
                  Statement('v2', '=', 'foo(not b and True)', '', np.float32)
                  ]
    scalar, vector = optimise_statements([], statements, variables)
    assert len(scalar) == 4
    assert scalar[0].expr == '1.0 * int(b and c)'
    assert scalar[1].expr == 'b and c'
    assert scalar[2].expr == 'int((not b) and True)'
    assert scalar[3].expr == '(not b) and True'
    assert len(vector) == 4
    assert vector[0].expr == '_lio_1'
    assert vector[1].expr == 'foo(_lio_2)'
    assert vector[2].expr == '_lio_3'
    assert vector[3].expr == 'foo(_lio_4)'

@pytest.mark.codegen_independent
def test_apply_loop_invariant_optimisation_no_optimisation():
    variables = {'v1': Variable('v1', scalar=False),
                 'v2': Variable('v2', scalar=False),
                 'N': Constant('N', 10),
                 's1': Variable('s1', scalar=True, dtype=float),
                 's2': Variable('s2', scalar=True, dtype=float),
                 'rand': DEFAULT_FUNCTIONS['rand']
                 }
    statements = [
        # This should not be simplified to 0!
        Statement('v1', '=', 'rand() - rand()', '', float),
        Statement('v1', '=', '3*rand() - 3*rand()', '', float),
        Statement('v1', '=', '3*rand() - ((1+2)*rand())', '', float),
        # This should not pull out rand()*N
        Statement('v1', '=', 's1*rand()*N', '', float),
        Statement('v1', '=', 's2*rand()*N', '', float),
        # This is not important mathematically, but it would change the numbers
        # that are generated
        Statement('v1', '=', '0*rand()*N', '', float),
        Statement('v1', '=', '0/rand()*N', '', float)
    ]
    scalar, vector = optimise_statements([], statements, variables)
    for vs in vector[:3]:
        assert vs.expr.count('rand()') == 2, f"Expression should still contain two rand() calls, but got {str(vs)}"
    for vs in vector[3:]:
        assert vs.expr.count('rand()') == 1, f"Expression should still contain a rand() call, but got {str(vs)}"

@pytest.mark.codegen_independent
def test_apply_loop_invariant_optimisation_simplification():
    variables = {'v1': Variable('v1', scalar=False),
                 'v2': Variable('v2', scalar=False),
                 'i1': Variable('i1', scalar=False, dtype=int),
                 'N': Constant('N', 10)
                 }
    statements = [
        # Should be simplified to 0.0
        Statement('v1', '=', 'v1 - v1', '', float),
        Statement('v1', '=', 'N*v1 - N*v1', '', float),
        Statement('v1', '=', 'v1*N * 0', '', float),
        Statement('v1', '=', 'v1 * 0', '', float),
        Statement('v1', '=', 'v1 * 0.0', '', float),
        Statement('v1', '=', '0.0 / (v1*N)', '', float),
        # Should be simplified to 0
        Statement('i1', '=', 'i1*N * 0', '', int),
        Statement('i1', '=', '0 * i1', '', int),
        Statement('i1', '=', '0 * i1*N', '', int),
        Statement('i1', '=', 'i1 * 0', '', int),
        # Should be simplified to v1*N
        Statement('v2', '=', '0 + v1*N', '', float),
        Statement('v2', '=', 'v1*N + 0.0', '', float),
        Statement('v2', '=', 'v1*N - 0', '', float),
        Statement('v2', '=', 'v1*N - 0.0', '', float),
        Statement('v2', '=', '1 * v1*N', '', float),
        Statement('v2', '=', '1.0 * v1*N', '', float),
        Statement('v2', '=', 'v1*N / 1.0', '', float),
        Statement('v2', '=', 'v1*N / 1', '', float),
        # Should be simplified to i1
        Statement('i1', '=', 'i1*1', '', int),
        Statement('i1', '=', 'i1//1', '', int),
        Statement('i1', '=', 'i1+0', '', int),
        Statement('i1', '=', '0+i1', '', int),
        Statement('i1', '=', 'i1-0', '', int),
        # Should *not* be simplified (because it would change the type,
        # important for integer division, for example)
        Statement('v1', '=', 'i1*1.0', '', float),
        Statement('v1', '=', '1.0*i1', '', float),
        Statement('v1', '=', 'i1/1.0', '', float),
        Statement('v1', '=', 'i1/1', '', float),
        Statement('v1', '=', 'i1+0.0', '', float),
        Statement('v1', '=', '0.0+i1', '', float),
        Statement('v1', '=', 'i1-0.0', '', float),
        ## Should *not* be simplified, flooring division by 1 changes the value
        Statement('v1', '=', 'v2//1.0', '', float),
        Statement('i1', '=', 'i1//1.0', '', float)  # changes type
    ]
    scalar, vector = optimise_statements([], statements, variables)
    assert len(scalar) == 0
    for s in vector[:6]:
        assert s.expr == '0.0'
    for s in vector[6:10]:
        assert s.expr == '0', s.expr  # integer
    for s in vector[10:18]:
        expr = s.expr.replace(' ', '')
        assert expr == 'v1*N' or expr == 'N*v1'
    for s in vector[18:23]:
        expr = s.expr.replace(' ', '')
        assert expr == 'i1'
    for s in vector[23:27]:
        expr = s.expr.replace(' ', '')
        assert expr == '1.0*i1' or expr == 'i1*1.0' or expr == 'i1/1.0'
    for s in vector[27:30]:
        expr = s.expr.replace(' ', '')
        assert expr == '0.0+i1' or expr == 'i1+0.0'
    for s in vector[30:31]:
        expr = s.expr.replace(' ', '')
        assert expr == 'v2//1.0' or expr == 'v2//1'
    for s in vector[31:]:
        expr = s.expr.replace(' ', '')
        assert expr == 'i1//1.0'


@pytest.mark.codegen_independent
def test_apply_loop_invariant_optimisation_constant_evaluation():
    variables = {'v1': Variable('v1', scalar=False),
                 'v2': Variable('v2', scalar=False),
                 'i1': Variable('i1', scalar=False, dtype=int),
                 'N': Constant('N', 10),
                 's1': Variable('s1', scalar=True, dtype=float),
                 's2': Variable('s2', scalar=True, dtype=float),
                 'exp': DEFAULT_FUNCTIONS['exp']
                 }
    statements = [
        Statement('v1', '=', 'v1 * (1 + 2 + 3)', '', float),
        Statement('v1', '=', 'exp(N)*v1', '', float),
        Statement('v1', '=', 'exp(0)*v1', '', float),
    ]
    scalar, vector = optimise_statements([], statements, variables)
    # exp(N) should be pulled out of the vector statements, the rest should be
    # evaluated in place
    assert len(scalar) == 1
    assert scalar[0].expr == 'exp(N)'
    assert len(vector) == 3
    expr = vector[0].expr.replace(' ', '')
    assert expr == '_lio_1*v1' or 'v1*_lio_1'
    expr = vector[1].expr.replace(' ', '')
    assert expr == '6.0*v1' or 'v1*6.0'
    assert vector[2].expr == 'v1'


@pytest.mark.codegen_independent
def test_automatic_augmented_assignments():
    # We test that statements that could be rewritten as augmented assignments
    # are correctly rewritten (using sympy to test for symbolic equality)
    variables = {
        'x': ArrayVariable('x', owner=None, size=10,
                           device=device),
        'y': ArrayVariable('y', owner=None, size=10,
                           device=device),
        'z': ArrayVariable('y', owner=None, size=10,
                           device=device),
        'b': ArrayVariable('b', owner=None, size=10,
                           dtype=bool, device=device),
        'clip': DEFAULT_FUNCTIONS['clip'],
        'inf': DEFAULT_CONSTANTS['inf']
    }
    statements = [
        # examples that should be rewritten
        # Note that using our approach, we will never get -= or /= but always
        # the equivalent += or *= statements
        ('x = x + 1.0', 'x += 1.0'),
        ('x = 2.0 * x', 'x *= 2.0'),
        ('x = x - 3.0', 'x += -3.0'),
        ('x = x/2.0', 'x *= 0.5'),
        ('x = y + (x + 1.0)', 'x += y + 1.0'),
        ('x = x + x', 'x *= 2.0'),
        ('x = x + y + z', 'x += y + z'),
        ('x = x + y + z', 'x += y + z'),
        # examples that should not be rewritten
        ('x = 1.0/x', 'x = 1.0/x'),
        ('x = 1.0', 'x = 1.0'),
        ('x = 2.0*(x + 1.0)', 'x = 2.0*(x + 1.0)'),
        ('x = clip(x + y, 0.0, inf)', 'x = clip(x + y, 0.0, inf)'),
        ('b = b or False', 'b = b or False')
    ]
    for orig, rewritten in statements:
        scalar, vector = make_statements(orig, variables, np.float32)
        try:  # we augment the assertion error with the original statement
            assert len(scalar) == 0, f"Did not expect any scalar statements but got {str(scalar)}"
            assert len(vector) == 1, f"Did expect a single statement but got {str(vector)}"
            statement = vector[0]
            expected_var, expected_op, expected_expr, _ = parse_statement(rewritten)
            assert expected_var == statement.var, f'expected write to variable {expected_var}, not to {statement.var}'
            assert expected_op == statement.op, f'expected operation {expected_op}, not {statement.op}'
            # Compare the two expressions using sympy to allow for different order etc.
            sympy_expected = str_to_sympy(expected_expr)
            sympy_actual = str_to_sympy(statement.expr)
            assert sympy_expected == sympy_actual, (
                f"RHS expressions '{sympy_to_str(sympy_expected)}' and '{sympy_to_str(sympy_actual)}' are not identical")
        except AssertionError as ex:
            raise AssertionError(
                f"Transformation for statement '{orig}' gave an unexpected result: {ex}")


def test_clear_cache():
    target = prefs.codegen.target
    if target == 'numpy':
        assert 'numpy' not in _cache_dirs_and_extensions
        with pytest.raises(ValueError):
            clear_cache('numpy')
    else:
        assert target in _cache_dirs_and_extensions
        cache_dir, _ = _cache_dirs_and_extensions[target]
        # Create a file that should not be there
        fname = os.path.join(cache_dir, 'some_file.py')
        open(fname, 'w').close()
        # clear_cache should refuse to clear the directory
        with pytest.raises(IOError):
            clear_cache(target)

        os.remove(fname)

@pytest.mark.skipif(platform.system() == 'Windows',
                    reason='CC and CXX variables are ignored on Windows.')
def test_compiler_error():
    # In particular on OSX with clang in a conda environment, compilation might fail.
    # Switching to a system gcc might help in such cases. Make sure that the error
    # message mentions that.
    old_CC = os.environ.get('CC', None)
    old_CXX = os.environ.get('CXX', None)
    os.environ.update({'CC': 'non-existing-compiler',
                       'CXX': 'non-existing-compiler++'})
    try:
        with catch_logs() as l:
            assert not CythonCodeObject.is_available()
        assert len(l) > 0  # There are additional warnings about compiler flags
        last_warning = l[-1]
        assert last_warning[1].endswith('.failed_compile_test')
        assert 'CC' in last_warning[2] and 'CXX' in last_warning[2]

    finally:
        if old_CC:
            os.environ['CC'] = old_CC
        else:
            del os.environ['CC']
        if old_CXX:
            os.environ['CXX'] = old_CXX
        else:
            del os.environ['CXX']


def test_compiler_c99():
    # On a user's computer, we do not know whether the compiler actually
    # has C99 support, so we just check whether the test does not raise an
    # error

    # The compiler check previously created spurious '-.o' files (see #1348)
    if os.path.exists('-.o'):
        os.remove('-.o')
    c99_support = compiler_supports_c99()
    assert not os.path.exists('-.o')
    # On our Azure test server we know that the compilers support C99
    if os.environ.get('AGENT_OS', ''):
        assert c99_support


def test_cpp_flags_support():
    from distutils.ccompiler import get_default_compiler
    from brian2.codegen.cpp_prefs import _compiler_flag_compatibility
    _compiler_flag_compatibility.clear()  # make sure cache is empty
    compiler = get_default_compiler()
    if compiler == 'msvc':
        pytest.skip('No flag support check for msvc')
    old_prefs = prefs['codegen.cpp.extra_compile_args']

    # Should always be supported
    prefs['codegen.cpp.extra_compile_args'] = ['-w']
    _, compile_args = get_compiler_and_args()
    assert compile_args == prefs['codegen.cpp.extra_compile_args']

    # Should never be supported and raise a warning
    prefs['codegen.cpp.extra_compile_args'] = ['-invalidxyz']
    with catch_logs() as l:
        _, compile_args = get_compiler_and_args()
    assert len(l) == 1 and l[0][0] == 'WARNING'
    assert compile_args == []

    prefs['codegen.cpp.extra_compile_args'] = old_prefs


@pytest.mark.skipif(platform.system() != 'Windows',
                    reason='MSVC flags are only relevant on Windows')
@pytest.mark.skipif(prefs['codegen.target'] == 'numpy',
                    reason='Test only relevant for compiled code')
def test_msvc_flags():
    # Very basic test that flags are stored to disk
    import brian2.codegen.cpp_prefs as cpp_prefs
    user_dir = os.path.join(os.path.expanduser('~'), '.brian')
    flag_file = os.path.join(user_dir, 'cpu_flags.txt')
    assert len(cpp_prefs.msvc_arch_flag)
    assert os.path.exists(flag_file)
    with open(flag_file, 'r', encoding='utf-8') as f:
        previously_stored_flags = json.load(f)
    hostname = socket.gethostname()
    assert hostname in previously_stored_flags
    assert len(previously_stored_flags[hostname])


if __name__ == '__main__':
    test_auto_target()
    test_analyse_identifiers()
    test_get_identifiers_recursively()
    test_write_to_subexpression()
    test_repeated_subexpressions()
    test_nested_subexpressions()
    test_apply_loop_invariant_optimisation()
    test_apply_loop_invariant_optimisation_integer()
    test_apply_loop_invariant_optimisation_boolean()
    test_apply_loop_invariant_optimisation_no_optimisation()
    test_apply_loop_invariant_optimisation_simplification()
    test_apply_loop_invariant_optimisation_constant_evaluation()
    test_automatic_augmented_assignments()
    test_clear_cache()
    test_msvc_flags()
