"""Tests for the basic tensor facilities using free algebra."""

import types

import pytest
from pyspark import SparkContext, SparkConf
from sympy import sympify, IndexedBase, sin, cos, KroneckerDelta, symbols

from drudge import Drudge, Range, Vec, Term


@pytest.fixture(scope='module')
def free_alg():
    """Initialize the environment for a free algebra."""

    conf = SparkConf().setMaster('local[2]').setAppName('free-alegra')
    ctx = SparkContext(conf=conf)
    dr = Drudge(ctx)

    r = Range('R')
    dumms = sympify('i, j, k, l, m, n')
    dr.set_dumms(r, dumms)

    s = Range('S')
    s_dumms = symbols('alpha beta')
    dr.set_dumms(s, s_dumms)

    v = Vec('v')

    return dr, types.SimpleNamespace(
        r=r, dumms=dumms, s=s, s_dumms=s_dumms, v=v
    )


def test_tensor_can_be_created(free_alg):
    """Test simple tensor creation."""

    dr, p = free_alg
    i = p.dumms[0]
    x = IndexedBase('x')
    tensor = dr.sum((i, p.r), x[i] * p.v[i])

    assert tensor.n_terms == 1

    terms = tensor.local_terms
    assert len(terms) == 1
    term = terms[0]
    assert term == Term([(i, p.r)], x[i], [p.v[i]])


def test_tensor_has_basic_operations(free_alg):
    """Test some of the basic operations on tensors.

    Tested in this module:

        1. Addition.
        2. Merge.
        3. Free variable.
        4. Dummy reset.
        5. Equality comparison.
    """

    dr, p = free_alg
    i, j, k, l, m = p.dumms[:5]
    x = IndexedBase('x')
    r = p.r
    v = p.v
    tensor = (
        dr.sum((l, r), x[i, l] * v[l]) +
        dr.sum((m, r), x[j, m] * v[m])
    )

    # Without dummy resetting, they cannot be merged.
    assert tensor.n_terms == 2
    assert tensor.merge().n_terms == 2

    # Free variables are important for dummy resetting.
    free_vars = tensor.free_vars
    assert free_vars == {x.label, i, j}

    # Reset dummy.
    reset = tensor.reset_dumms()
    expected = (
        dr.sum((k, r), x[i, k] * v[k]) +
        dr.sum((k, r), x[j, k] * v[k])
    )
    assert reset == expected
    assert reset.local_terms == expected.local_terms

    # Merge the terms.
    merged = reset.merge()
    assert merged.n_terms == 1
    term = merged.local_terms[0]
    assert term == Term([(k, r)], x[i, k] + x[j, k], [v[k]])

@@ -2257,12 +2257,1255 @@
 k%5D, %5Bv%5Bk%5D%5D)%0A
+%0A%0Adef test_tensor_can_be_simplified_amp(free_alg):%0A    %22%22%22Test the amplitude simplification for tensors.%0A%0A    More than trivial tensor amplitude simplification is tested here.  Currently%0A    it mostly concentrates on the dispatching to SymPy and delta simplification.%0A    The master simplification is also tested.%0A    %22%22%22%0A%0A    dr, p = free_alg%0A    r = p.r%0A    s = p.s%0A    v = p.v%0A    i, j = p.dumms%5B:2%5D%0A    alpha = p.s_dumms%5B0%5D%0A%0A    x = IndexedBase('x')%0A    y = IndexedBase('y')%0A    theta = sympify('theta')%0A%0A    tensor = (%0A        dr.sum((i, r), sin(theta) ** 2 * x%5Bi%5D * v%5Bi%5D) +%0A        dr.sum(%0A            (i, r), (j, r),%0A            cos(theta) ** 2 * x%5Bj%5D * KroneckerDelta(i, j) * v%5Bi%5D%0A        ) +%0A        dr.sum((i, r), (alpha, s), KroneckerDelta(i, alpha) * y%5Bi%5D * v%5Bi%5D)%0A    )%0A    assert tensor.n_terms == 3%0A%0A    first = tensor.simplify_amps()%0A    # Now we should have one term killed.%0A    assert first.n_terms == 2%0A%0A    # Merge again should really simplify.%0A    merged = first.reset_dumms().merge().simplify_amps()%0A    assert merged.n_terms == 1%0A    expected = dr.sum((i, r), x%5Bi%5D * v%5Bi%5D)%0A    assert merged == expected%0A%0A    # The master simplification should do it in one turn.%0A    simpl = tensor.simplify()%0A    assert simpl == expected%0A

Add tests for tensor amplitude simplification