from tests._helpers import (
  unittest, parameterized, Fraction, time,
  run_test, run_compiler_error_test,
  ArrayIndex, ArraySlice, Bool, BinOp, BuiltIn, DataArray, DataDict, DataRangeArray,
  Defini, Continva, Designa, DesignaDestructure, DesignaIndex, DumStatement,
  Erumpe, ExpressionStatement, Fvnctio, ID, InterpolatedString, Invoca,
  ModuleCall, Nullus, Numeral, PerStatement, Program, Redi, SiStatement,
  String, TemptaStatement, UnaryMinus, UnaryNot, Fractio, frac_to_fraction,
  fraction_to_frac, num_to_int, int_to_num, make_string,
  ValInt, ValStr, ValBool, ValList, ValDict, ValNul, ValFunc, ValFrac,
  CentvrionError, _RUNTIME_C, _cent_rng,
  Lexer, Parser, compile_program,
  os, subprocess, tempfile, StringIO, patch,
)


# --- Array concatenation ---

array_concat_tests = [
  ("[I, II] @ [III, IV]",
    Program([], [ExpressionStatement(BinOp(DataArray([Numeral("I"), Numeral("II")]), DataArray([Numeral("III"), Numeral("IV")]), "SYMBOL_AT"))]),
    ValList([ValInt(1), ValInt(2), ValInt(3), ValInt(4)])),
  ("[] @ [I]",
    Program([], [ExpressionStatement(BinOp(DataArray([]), DataArray([Numeral("I")]), "SYMBOL_AT"))]),
    ValList([ValInt(1)])),
  ("[I] @ []",
    Program([], [ExpressionStatement(BinOp(DataArray([Numeral("I")]), DataArray([]), "SYMBOL_AT"))]),
    ValList([ValInt(1)])),
  ("[] @ []",
    Program([], [ExpressionStatement(BinOp(DataArray([]), DataArray([]), "SYMBOL_AT"))]),
    ValList([])),
  ('["a"] @ [I]',
    Program([], [ExpressionStatement(BinOp(DataArray([String("a")]), DataArray([Numeral("I")]), "SYMBOL_AT"))]),
    ValList([ValStr("a"), ValInt(1)])),
  # left-associative chaining
  ("[I] @ [II] @ [III]",
    Program([], [ExpressionStatement(BinOp(BinOp(DataArray([Numeral("I")]), DataArray([Numeral("II")]), "SYMBOL_AT"), DataArray([Numeral("III")]), "SYMBOL_AT"))]),
    ValList([ValInt(1), ValInt(2), ValInt(3)])),
  # concat with variable
  ("DESIGNA a VT [I, II]\nDESIGNA b VT [III]\na @ b",
    Program([], [Designa(ID("a"), DataArray([Numeral("I"), Numeral("II")])), Designa(ID("b"), DataArray([Numeral("III")])), ExpressionStatement(BinOp(ID("a"), ID("b"), "SYMBOL_AT"))]),
    ValList([ValInt(1), ValInt(2), ValInt(3)])),
]

class TestArrayConcat(unittest.TestCase):
  @parameterized.expand(array_concat_tests)
  def test_array_concat(self, source, nodes, value):
    run_test(self, source, nodes, value)

# --- Array indexing ---
# Indexing is 1-based; I is the first element

array_index_tests = [
  # basic indexing
  ("[I, II, III][I]", Program([], [ExpressionStatement(ArrayIndex(DataArray([Numeral("I"), Numeral("II"), Numeral("III")]), Numeral("I")))]), ValInt(1)),    # first element
  ("[I, II, III][II]", Program([], [ExpressionStatement(ArrayIndex(DataArray([Numeral("I"), Numeral("II"), Numeral("III")]), Numeral("II")))]), ValInt(2)),   # second element
  ("[I, II, III][III]", Program([], [ExpressionStatement(ArrayIndex(DataArray([Numeral("I"), Numeral("II"), Numeral("III")]), Numeral("III")))]), ValInt(3)),  # third element
  # index into a variable
  ("DESIGNA a VT [X, XX, XXX]\na[II]",
    Program([], [Designa(ID("a"), DataArray([Numeral("X"), Numeral("XX"), Numeral("XXX")])), ExpressionStatement(ArrayIndex(ID("a"), Numeral("II")))]),
    ValInt(20)),  # second element
  # index into range array
  ("[I VSQVE V][II]", Program([], [ExpressionStatement(ArrayIndex(DataRangeArray(Numeral("I"), Numeral("V")), Numeral("II")))]), ValInt(2)),      # second element of [1,2,3,4,5]
  # expression as index
  ("[I, II, III][I + I]",
    Program([], [ExpressionStatement(ArrayIndex(
      DataArray([Numeral("I"), Numeral("II"), Numeral("III")]),
      BinOp(Numeral("I"), Numeral("I"), "SYMBOL_PLUS")))]),
    ValInt(2)),
  # division result as index (no FRACTIO): IV / II = 2
  ("[X, XX, XXX][IV / II]",
    Program([], [ExpressionStatement(ArrayIndex(
      DataArray([Numeral("X"), Numeral("XX"), Numeral("XXX")]),
      BinOp(Numeral("IV"), Numeral("II"), "SYMBOL_DIVIDE")))]),
    ValInt(20)),
  # whole-number fraction (from division) as index, with FRACTIO imported
  ("CVM FRACTIO\n[X, XX, XXX][IV / II]",
    Program([ModuleCall("FRACTIO")], [ExpressionStatement(ArrayIndex(
      DataArray([Numeral("X"), Numeral("XX"), Numeral("XXX")]),
      BinOp(Numeral("IV"), Numeral("II"), "SYMBOL_DIVIDE")))]),
    ValInt(20)),
]

class TestArrayIndex(unittest.TestCase):
  @parameterized.expand(array_index_tests)
  def test_array_index(self, source, nodes, value):
    run_test(self, source, nodes, value)

# --- Array index assignment ---

array_index_assign_tests = [
  # assign to middle element
  ("DESIGNA a VT [I, II, III]\nDESIGNA a[II] VT X\na[II]",
    Program([], [
      Designa(ID("a"), DataArray([Numeral("I"), Numeral("II"), Numeral("III")])),
      DesignaIndex(ID("a"), [Numeral("II")], Numeral("X")),
      ExpressionStatement(ArrayIndex(ID("a"), Numeral("II"))),
    ]),
    ValInt(10)),
  # assign to first element
  ("DESIGNA a VT [I, II, III]\nDESIGNA a[I] VT V\na[I]",
    Program([], [
      Designa(ID("a"), DataArray([Numeral("I"), Numeral("II"), Numeral("III")])),
      DesignaIndex(ID("a"), [Numeral("I")], Numeral("V")),
      ExpressionStatement(ArrayIndex(ID("a"), Numeral("I"))),
    ]),
    ValInt(5)),
  # assign to last element
  ("DESIGNA a VT [I, II, III]\nDESIGNA a[III] VT L\na[III]",
    Program([], [
      Designa(ID("a"), DataArray([Numeral("I"), Numeral("II"), Numeral("III")])),
      DesignaIndex(ID("a"), [Numeral("III")], Numeral("L")),
      ExpressionStatement(ArrayIndex(ID("a"), Numeral("III"))),
    ]),
    ValInt(50)),
  # other elements unaffected
  ("DESIGNA a VT [I, II, III]\nDESIGNA a[II] VT X\na[I]",
    Program([], [
      Designa(ID("a"), DataArray([Numeral("I"), Numeral("II"), Numeral("III")])),
      DesignaIndex(ID("a"), [Numeral("II")], Numeral("X")),
      ExpressionStatement(ArrayIndex(ID("a"), Numeral("I"))),
    ]),
    ValInt(1)),
  # expression as index
  ("DESIGNA a VT [I, II, III]\nDESIGNA i VT II\nDESIGNA a[i] VT X\na[II]",
    Program([], [
      Designa(ID("a"), DataArray([Numeral("I"), Numeral("II"), Numeral("III")])),
      Designa(ID("i"), Numeral("II")),
      DesignaIndex(ID("a"), [ID("i")], Numeral("X")),
      ExpressionStatement(ArrayIndex(ID("a"), Numeral("II"))),
    ]),
    ValInt(10)),
]

class TestArrayIndexAssign(unittest.TestCase):
  @parameterized.expand(array_index_assign_tests)
  def test_array_index_assign(self, source, nodes, value):
    run_test(self, source, nodes, value)

# --- Multi-dimensional array index assignment ---

multidim_assign_tests = [
  # 2D array assignment
  ("DESIGNA a VT [[I, II], [III, IV]]\nDESIGNA a[I][II] VT X\na[I][II]",
    Program([], [
      Designa(ID("a"), DataArray([DataArray([Numeral("I"), Numeral("II")]), DataArray([Numeral("III"), Numeral("IV")])])),
      DesignaIndex(ID("a"), [Numeral("I"), Numeral("II")], Numeral("X")),
      ExpressionStatement(ArrayIndex(ArrayIndex(ID("a"), Numeral("I")), Numeral("II"))),
    ]),
    ValInt(10)),
  # other elements unaffected
  ("DESIGNA a VT [[I, II], [III, IV]]\nDESIGNA a[I][II] VT X\na[II][I]",
    Program([], [
      Designa(ID("a"), DataArray([DataArray([Numeral("I"), Numeral("II")]), DataArray([Numeral("III"), Numeral("IV")])])),
      DesignaIndex(ID("a"), [Numeral("I"), Numeral("II")], Numeral("X")),
      ExpressionStatement(ArrayIndex(ArrayIndex(ID("a"), Numeral("II")), Numeral("I"))),
    ]),
    ValInt(3)),
  # dict inside array
  ('DESIGNA a VT [TABVLA {"x" VT I}]\nDESIGNA a[I]["x"] VT X\na[I]["x"]',
    Program([], [
      Designa(ID("a"), DataArray([DataDict([(String("x"), Numeral("I"))])])),
      DesignaIndex(ID("a"), [Numeral("I"), String("x")], Numeral("X")),
      ExpressionStatement(ArrayIndex(ArrayIndex(ID("a"), Numeral("I")), String("x"))),
    ]),
    ValInt(10)),
  # array inside dict
  ('DESIGNA d VT TABVLA {"a" VT [I, II]}\nDESIGNA d["a"][I] VT X\nd["a"][I]',
    Program([], [
      Designa(ID("d"), DataDict([(String("a"), DataArray([Numeral("I"), Numeral("II")]))])),
      DesignaIndex(ID("d"), [String("a"), Numeral("I")], Numeral("X")),
      ExpressionStatement(ArrayIndex(ArrayIndex(ID("d"), String("a")), Numeral("I"))),
    ]),
    ValInt(10)),
  # 3 levels deep
  ("DESIGNA a VT [[[I]]]\nDESIGNA a[I][I][I] VT X\na[I][I][I]",
    Program([], [
      Designa(ID("a"), DataArray([DataArray([DataArray([Numeral("I")])])])),
      DesignaIndex(ID("a"), [Numeral("I"), Numeral("I"), Numeral("I")], Numeral("X")),
      ExpressionStatement(ArrayIndex(ArrayIndex(ArrayIndex(ID("a"), Numeral("I")), Numeral("I")), Numeral("I"))),
    ]),
    ValInt(10)),
]

class TestMultidimAssign(unittest.TestCase):
  @parameterized.expand(multidim_assign_tests)
  def test_multidim_assign(self, source, nodes, value):
    run_test(self, source, nodes, value)

# --- Array slicing ---

array_slice_tests = [
  # basic slice from middle
  ("[X, XX, XXX, XL, L][II VSQVE IV]",
    Program([], [ExpressionStatement(ArraySlice(
      DataArray([Numeral("X"), Numeral("XX"), Numeral("XXX"), Numeral("XL"), Numeral("L")]),
      Numeral("II"), Numeral("IV")))]),
    ValList([ValInt(20), ValInt(30), ValInt(40)])),
  # slice of length 1
  ("[I, II, III][II VSQVE II]",
    Program([], [ExpressionStatement(ArraySlice(
      DataArray([Numeral("I"), Numeral("II"), Numeral("III")]),
      Numeral("II"), Numeral("II")))]),
    ValList([ValInt(2)])),
  # full array slice
  ("[I, II, III][I VSQVE III]",
    Program([], [ExpressionStatement(ArraySlice(
      DataArray([Numeral("I"), Numeral("II"), Numeral("III")]),
      Numeral("I"), Numeral("III")))]),
    ValList([ValInt(1), ValInt(2), ValInt(3)])),
  # slice on variable
  ("DESIGNA a VT [I, II, III, IV, V]\na[II VSQVE IV]",
    Program([], [
      Designa(ID("a"), DataArray([Numeral("I"), Numeral("II"), Numeral("III"), Numeral("IV"), Numeral("V")])),
      ExpressionStatement(ArraySlice(ID("a"), Numeral("II"), Numeral("IV"))),
    ]),
    ValList([ValInt(2), ValInt(3), ValInt(4)])),
  # slice then index (chained)
  ("[I, II, III, IV][I VSQVE III][II]",
    Program([], [ExpressionStatement(ArrayIndex(
      ArraySlice(
        DataArray([Numeral("I"), Numeral("II"), Numeral("III"), Numeral("IV")]),
        Numeral("I"), Numeral("III")),
      Numeral("II")))]),
    ValInt(2)),
  # slice on range array
  ("[I VSQVE X][III VSQVE VII]",
    Program([], [ExpressionStatement(ArraySlice(
      DataRangeArray(Numeral("I"), Numeral("X")),
      Numeral("III"), Numeral("VII")))]),
    ValList([ValInt(3), ValInt(4), ValInt(5), ValInt(6), ValInt(7)])),
  # expression as slice bounds
  ("[I, II, III, IV, V][I + I VSQVE II + II]",
    Program([], [ExpressionStatement(ArraySlice(
      DataArray([Numeral("I"), Numeral("II"), Numeral("III"), Numeral("IV"), Numeral("V")]),
      BinOp(Numeral("I"), Numeral("I"), "SYMBOL_PLUS"),
      BinOp(Numeral("II"), Numeral("II"), "SYMBOL_PLUS")))]),
    ValList([ValInt(2), ValInt(3), ValInt(4)])),
]

class TestArraySlice(unittest.TestCase):
  @parameterized.expand(array_slice_tests)
  def test_array_slice(self, source, nodes, value):
    run_test(self, source, nodes, value)

# --- Multiline arrays ---

multiline_array_tests = [
  # newlines after commas
  ("[I,\nII,\nIII]",
    Program([], [ExpressionStatement(DataArray([Numeral("I"), Numeral("II"), Numeral("III")]))]),
    ValList([ValInt(1), ValInt(2), ValInt(3)])),
  # single newline after comma
  ("[I, II,\nIII]",
    Program([], [ExpressionStatement(DataArray([Numeral("I"), Numeral("II"), Numeral("III")]))]),
    ValList([ValInt(1), ValInt(2), ValInt(3)])),
  # empty array still works
  ("[]",
    Program([], [ExpressionStatement(DataArray([]))]),
    ValList([])),
  # nested arrays with newlines
  ("[I,\n[II, III],\nIV]",
    Program([], [ExpressionStatement(DataArray([
      Numeral("I"),
      DataArray([Numeral("II"), Numeral("III")]),
      Numeral("IV")]))]),
    ValList([ValInt(1), ValList([ValInt(2), ValInt(3)]), ValInt(4)])),
]

class TestMultilineArray(unittest.TestCase):
  @parameterized.expand(multiline_array_tests)
  def test_multiline_array(self, source, nodes, value):
    run_test(self, source, nodes, value)