from io import StringIO

import pytest

import diofant
from diofant import Equality, Piecewise, cos, sin, tan
from diofant.abc import A, B, C, a, t, x, y, z
from diofant.core import Catalan, Eq, EulerGamma, Function, pi, symbols, zoo
from diofant.matrices import Matrix, MatrixSymbol
from diofant.tensor import Idx, IndexedBase
from diofant.utilities.codegen import OctaveCodeGen, codegen, make_routine


__all__ = ()


def test_empty_m_code():
    code_gen = OctaveCodeGen()
    output = StringIO()
    code_gen.dump_m([], output, "file", header=False, empty=False)
    source = output.getvalue()
    assert source == ""


def test_m_simple_code():
    name_expr = ("test", (x + y)*z)
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    assert result[0] == "test.m"
    source = result[1]
    expected = (
        "function out1 = test(x, y, z)\n"
        "  out1 = z.*(x + y);\n"
        "end\n"
    )
    assert source == expected


def test_m_simple_code_with_header():
    name_expr = ("test", (x + y)*z)
    result, = codegen(name_expr, "Octave", header=True, empty=False)
    assert result[0] == "test.m"
    source = result[1]
    expected = (
        "function out1 = test(x, y, z)\n"
        "  %TEST  Autogenerated by diofant\n"
        "  %   Code generated with diofant " + diofant.__version__ + "\n"
        "  %\n"
        "  %   See https://diofant.readthedocs.io/ for more information.\n"
        "  %\n"
        "  %   This file is part of 'project'\n"
        "  out1 = z.*(x + y);\n"
        "end\n"
    )
    assert source == expected


def test_m_simple_code_nameout():
    expr = Equality(z, (x + y))
    name_expr = ("test", expr)
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function z = test(x, y)\n"
        "  z = x + y;\n"
        "end\n"
    )
    assert source == expected


def test_m_numbersymbol():
    name_expr = ("test", pi**Catalan)
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function out1 = test()\n"
        "  out1 = pi^0.915965594177219;\n"
        "end\n"
    )
    assert source == expected


@pytest.mark.xfail
def test_m_numbersymbol_no_inline():
    # FIXME: how to pass inline=False to the OctaveCodePrinter?
    name_expr = ("test", [pi**Catalan, EulerGamma])
    result, = codegen(name_expr, "Octave", header=False,
                      empty=False, inline=False)
    source = result[1]
    expected = (
        "function [out1, out2] = test()\n"
        "  Catalan = 0.915965594177219;  % constant\n"
        "  EulerGamma = 0.5772156649015329;  % constant\n"
        "  out1 = pi^Catalan;\n"
        "  out2 = EulerGamma;\n"
        "end\n"
    )
    assert source == expected


def test_m_code_argument_order():
    expr = x + y
    routine = make_routine("test", expr, argument_sequence=[z, x, y], language="octave")
    code_gen = OctaveCodeGen()
    output = StringIO()
    code_gen.dump_m([routine], output, "test", header=False, empty=False)
    source = output.getvalue()
    expected = (
        "function out1 = test(z, x, y)\n"
        "  out1 = x + y;\n"
        "end\n"
    )
    assert source == expected


def test_multiple_results_m():
    # Here the output order is the input order
    expr1 = (x + y)*z
    expr2 = (x - y)*z
    name_expr = ("test", [expr1, expr2])
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function [out1, out2] = test(x, y, z)\n"
        "  out1 = z.*(x + y);\n"
        "  out2 = z.*(x - y);\n"
        "end\n"
    )
    assert source == expected


def test_results_named_unordered():
    # Here output order is based on name_expr
    expr1 = Equality(C, (x + y)*z)
    expr2 = Equality(A, (x - y)*z)
    expr3 = Equality(B, 2*x)
    name_expr = ("test", [expr1, expr2, expr3])
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function [C, A, B] = test(x, y, z)\n"
        "  C = z.*(x + y);\n"
        "  A = z.*(x - y);\n"
        "  B = 2*x;\n"
        "end\n"
    )
    assert source == expected


def test_results_named_ordered():
    expr1 = Equality(C, (x + y)*z)
    expr2 = Equality(A, (x - y)*z)
    expr3 = Equality(B, 2*x)
    name_expr = ("test", [expr1, expr2, expr3])
    result = codegen(name_expr, "Octave", header=False, empty=False,
                     argument_sequence=(x, z, y))
    assert result[0][0] == "test.m"
    source = result[0][1]
    expected = (
        "function [C, A, B] = test(x, z, y)\n"
        "  C = z.*(x + y);\n"
        "  A = z.*(x - y);\n"
        "  B = 2*x;\n"
        "end\n"
    )
    assert source == expected


def test_complicated_m_codegen():
    name_expr = ("testlong",
                 [((sin(x) + cos(y) + tan(z))**3).expand(),
                  cos(cos(cos(cos(cos(cos(cos(cos(x + y + z))))))))])
    result = codegen(name_expr, "Octave", header=False, empty=False)
    assert result[0][0] == "testlong.m"
    source = result[0][1]
    expected = (
        "function [out1, out2] = testlong(x, y, z)\n"
        "  out1 = sin(x).^3 + 3*sin(x).^2.*cos(y) + 3*sin(x).^2.*tan(z)"
        " + 3*sin(x).*cos(y).^2 + 6*sin(x).*cos(y).*tan(z) + 3*sin(x).*tan(z).^2"
        " + cos(y).^3 + 3*cos(y).^2.*tan(z) + 3*cos(y).*tan(z).^2 + tan(z).^3;\n"
        "  out2 = cos(cos(cos(cos(cos(cos(cos(cos(x + y + z))))))));\n"
        "end\n"
    )
    assert source == expected


def test_m_output_arg_mixed_unordered():
    # named outputs are alphabetical, unnamed output appear in the given order
    name_expr = ("foo", [cos(2*x), Equality(y, sin(x)), cos(x), Equality(a, sin(2*x))])
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    assert result[0] == "foo.m"
    source = result[1]
    expected = (
        'function [out1, y, out3, a] = foo(x)\n'
        '  out1 = cos(2*x);\n'
        '  y = sin(x);\n'
        '  out3 = cos(x);\n'
        '  a = sin(2*x);\n'
        'end\n'
    )
    assert source == expected


def test_m_piecewise_():
    pw = Piecewise((0, x < -1), (x**2, x <= 1), (-x+2, x > 1), (1, True))
    name_expr = ("pwtest", pw)
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function out1 = pwtest(x)\n"
        "  out1 = ((x < -1).*(0) + (~(x < -1)).*( ...\n"
        "  (x <= 1).*(x.^2) + (~(x <= 1)).*( ...\n"
        "  (x > 1).*(-x + 2) + (~(x > 1)).*(1))));\n"
        "end\n"
    )
    assert source == expected


@pytest.mark.xfail
def test_m_piecewise_no_inline():
    # FIXME: how to pass inline=False to the OctaveCodePrinter?
    pw = Piecewise((0, x < -1), (x**2, x <= 1), (-x+2, x > 1), (1, True))
    name_expr = ("pwtest", pw)
    result, = codegen(name_expr, "Octave", header=False, empty=False,
                      inline=False)
    source = result[1]
    expected = (
        "function out1 = pwtest(x)\n"
        "  if (x < -1)\n"
        "    out1 = 0;\n"
        "  elseif (x <= 1)\n"
        "    out1 = x.^2;\n"
        "  elseif (x > 1)\n"
        "    out1 = -x + 2;\n"
        "  else\n"
        "    out1 = 1;\n"
        "  end\n"
        "end\n"
    )
    assert source == expected


def test_m_multifcns_per_file():
    name_expr = [ ("foo", [2*x, 3*y]), ("bar", [y**2, 4*y]) ]
    result = codegen(name_expr, "Octave", header=False, empty=False)
    assert result[0][0] == "foo.m"
    source = result[0][1]
    expected = (
        "function [out1, out2] = foo(x, y)\n"
        "  out1 = 2*x;\n"
        "  out2 = 3*y;\n"
        "end\n"
        "function [out1, out2] = bar(y)\n"
        "  out1 = y.^2;\n"
        "  out2 = 4*y;\n"
        "end\n"
    )
    assert source == expected


def test_m_multifcns_per_file_w_header():
    name_expr = [ ("foo", [2*x, 3*y]), ("bar", [y**2, 4*y]) ]
    result = codegen(name_expr, "Octave", header=True, empty=False)
    assert result[0][0] == "foo.m"
    source = result[0][1]
    expected = (
        "function [out1, out2] = foo(x, y)\n"
        "  %FOO  Autogenerated by diofant\n"
        "  %   Code generated with diofant " + diofant.__version__ + "\n"
        "  %\n"
        "  %   See https://diofant.readthedocs.io/ for more information.\n"
        "  %\n"
        "  %   This file is part of 'project'\n"
        "  out1 = 2*x;\n"
        "  out2 = 3*y;\n"
        "end\n"
        "function [out1, out2] = bar(y)\n"
        "  out1 = y.^2;\n"
        "  out2 = 4*y;\n"
        "end\n"
    )
    assert source == expected


def test_m_filename_match_first_fcn():
    name_expr = [ ("foo", [2*x, 3*y]), ("bar", [y**2, 4*y]) ]
    pytest.raises(ValueError,
                  lambda: codegen(name_expr,
                                  "Octave", prefix="bar", header=False, empty=False))


def test_m_matrix_named():
    e2 = Matrix([[x, 2*y, pi*z]])
    name_expr = ("test", Equality(MatrixSymbol('myout1', 1, 3), e2))
    result = codegen(name_expr, "Octave", header=False, empty=False)
    assert result[0][0] == "test.m"
    source = result[0][1]
    expected = (
        "function myout1 = test(x, y, z)\n"
        "  myout1 = [x 2*y pi*z];\n"
        "end\n"
    )
    assert source == expected


def test_m_matrix_named_matsym():
    myout1 = MatrixSymbol('myout1', 1, 3)
    e2 = Matrix([[x, 2*y, pi*z]])
    name_expr = ("test", Equality(myout1, e2, evaluate=False))
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function myout1 = test(x, y, z)\n"
        "  myout1 = [x 2*y pi*z];\n"
        "end\n"
    )
    assert source == expected


def test_m_matrix_output_autoname():
    expr = Matrix([[x, x+y, 3]])
    name_expr = ("test", expr)
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function out1 = test(x, y)\n"
        "  out1 = [x x + y 3];\n"
        "end\n"
    )
    assert source == expected


def test_m_matrix_output_autoname_2():
    e1 = (x + y)
    e2 = Matrix([[2*x, 2*y, 2*z]])
    e3 = Matrix([[x], [y], [z]])
    e4 = Matrix([[x, y], [z, 16]])
    name_expr = ("test", (e1, e2, e3, e4))
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function [out1, out2, out3, out4] = test(x, y, z)\n"
        "  out1 = x + y;\n"
        "  out2 = [2*x 2*y 2*z];\n"
        "  out3 = [x; y; z];\n"
        "  out4 = [x  y;\n"
        "  z 16];\n"
        "end\n"
    )
    assert source == expected


def test_m_results_matrix_named_ordered():
    A = MatrixSymbol('A', 1, 3)
    expr1 = Equality(C, (x + y)*z)
    expr2 = Equality(A, Matrix([[1, 2, x]]))
    expr3 = Equality(B, 2*x)
    name_expr = ("test", [expr1, expr2, expr3])
    result, = codegen(name_expr, "Octave", header=False, empty=False,
                      argument_sequence=(x, z, y))
    source = result[1]
    expected = (
        "function [C, A, B] = test(x, z, y)\n"
        "  C = z.*(x + y);\n"
        "  A = [1 2 x];\n"
        "  B = 2*x;\n"
        "end\n"
    )
    assert source == expected


def test_m_matrixsymbol_slice():
    A = MatrixSymbol('A', 2, 3)
    B = MatrixSymbol('B', 1, 3)
    C = MatrixSymbol('C', 1, 3)
    D = MatrixSymbol('D', 2, 1)
    name_expr = ("test", [Equality(B, A[0, :]),
                          Equality(C, A[1, :]),
                          Equality(D, A[:, 2])])
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function [B, C, D] = test(A)\n"
        "  B = A(1, :);\n"
        "  C = A(2, :);\n"
        "  D = A(:, 3);\n"
        "end\n"
    )
    assert source == expected


def test_m_matrixsymbol_slice2():
    A = MatrixSymbol('A', 3, 4)
    B = MatrixSymbol('B', 2, 2)
    C = MatrixSymbol('C', 2, 2)
    name_expr = ("test", [Equality(B, A[0:2, 0:2]),
                          Equality(C, A[0:2, 1:3])])
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function [B, C] = test(A)\n"
        "  B = A(1:2, 1:2);\n"
        "  C = A(1:2, 2:3);\n"
        "end\n"
    )
    assert source == expected


def test_m_matrixsymbol_slice3():
    A = MatrixSymbol('A', 8, 7)
    B = MatrixSymbol('B', 2, 2)
    C = MatrixSymbol('C', 4, 2)
    name_expr = ("test", [Equality(B, A[6:, 1::3]),
                          Equality(C, A[::2, ::3])])
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function [B, C] = test(A)\n"
        "  B = A(7:end, 2:3:end);\n"
        "  C = A(1:2:end, 1:3:end);\n"
        "end\n"
    )
    assert source == expected


def test_m_matrixsymbol_slice_autoname():
    A = MatrixSymbol('A', 2, 3)
    B = MatrixSymbol('B', 1, 3)
    name_expr = ("test", [Equality(B, A[0, :]), A[1, :], A[:, 0], A[:, 1]])
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function [B, out2, out3, out4] = test(A)\n"
        "  B = A(1, :);\n"
        "  out2 = A(2, :);\n"
        "  out3 = A(:, 1);\n"
        "  out4 = A(:, 2);\n"
        "end\n"
    )
    assert source == expected


def test_m_loops():
    # Note: an Octave programmer would probably vectorize this across one or
    # more dimensions.  Also, size(A) would be used rather than passing in m
    # and n.  Perhaps users would expect us to vectorize automatically here?
    # Or is it possible to represent such things using IndexedBase?
    n, m = symbols('n m', integer=True)
    A = IndexedBase('A')
    x = IndexedBase('x')
    y = IndexedBase('y')
    i = Idx('i', m)
    j = Idx('j', n)
    result, = codegen(('mat_vec_mult', Eq(y[i], A[i, j]*x[j])), "Octave",
                      header=False, empty=False)
    source = result[1]
    expected = (
        'function y = mat_vec_mult(A, m, n, x)\n'
        '  for i = 1:m\n'
        '    y(i) = 0;\n'
        '  end\n'
        '  for i = 1:m\n'
        '    for j = 1:n\n'
        '      y(i) = %(rhs)s + y(i);\n'
        '    end\n'
        '  end\n'
        'end\n'
    )
    assert (source == expected % {'rhs': 'A(%s, %s).*x(j)' % (i, j)} or
            source == expected % {'rhs': 'x(j).*A(%s, %s)' % (i, j)})


def test_m_tensor_loops_multiple_contractions():
    # see comments in previous test about vectorizing
    n, m, o, p = symbols('n m o p', integer=True)
    A = IndexedBase('A')
    B = IndexedBase('B')
    y = IndexedBase('y')
    i = Idx('i', m)
    j = Idx('j', n)
    k = Idx('k', o)
    l = Idx('l', p)
    result, = codegen(('tensorthing', Eq(y[i], B[j, k, l]*A[i, j, k, l])),
                      "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        'function y = tensorthing(A, B, m, n, o, p)\n'
        '  for i = 1:m\n'
        '    y(i) = 0;\n'
        '  end\n'
        '  for i = 1:m\n'
        '    for j = 1:n\n'
        '      for k = 1:o\n'
        '        for l = 1:p\n'
        '          y(i) = y(i) + B(j, k, l).*A(i, j, k, l);\n'
        '        end\n'
        '      end\n'
        '    end\n'
        '  end\n'
        'end\n'
    )
    assert source == expected


def test_m_InOutArgument():
    expr = Equality(x, x**2)
    name_expr = ("mysqr", expr)
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function x = mysqr(x)\n"
        "  x = x.^2;\n"
        "end\n"
    )
    assert source == expected


def test_m_InOutArgument_order():
    # can specify the order as (x, y)
    expr = Equality(x, x**2 + y)
    name_expr = ("test", expr)
    result, = codegen(name_expr, "Octave", header=False,
                      empty=False, argument_sequence=(x, y))
    source = result[1]
    expected = (
        "function x = test(x, y)\n"
        "  x = x.^2 + y;\n"
        "end\n"
    )
    assert source == expected
    # make sure it gives (x, y) not (y, x)
    expr = Equality(x, x**2 + y)
    name_expr = ("test", expr)
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function x = test(x, y)\n"
        "  x = x.^2 + y;\n"
        "end\n"
    )
    assert source == expected


def test_m_not_supported():
    f = Function('f')
    name_expr = ("test", [f(x).diff(x), zoo])
    result, = codegen(name_expr, "Octave", header=False, empty=False)
    source = result[1]
    expected = (
        "function [out1, out2] = test(x)\n"
        "  % unsupported: Derivative(f(x), x)\n"
        "  % unsupported: zoo\n"
        "  out1 = Derivative(f(x), x);\n"
        "  out2 = zoo;\n"
        "end\n"
    )
    assert source == expected


def test_global_vars_octave():
    result = codegen(('f', x*y), "Octave", header=False, empty=False,
                     global_vars=(y,))
    source = result[0][1]
    expected = (
        "function out1 = f(x)\n"
        "  global y\n"
        "  out1 = x.*y;\n"
        "end\n"
    )
    assert source == expected

    result = codegen(('f', x*y+z), "Octave", header=False, empty=False,
                     argument_sequence=(x, y), global_vars=(z, t))
    source = result[0][1]
    expected = (
        "function out1 = f(x, y)\n"
        "  global t z\n"
        "  out1 = x.*y + z;\n"
        "end\n"
    )
    assert source == expected
