import os

import pytest
from sage.all import *  # noqa: F401

load(os.path.join(os.path.dirname(__file__), "..", "sage", "period_computation.sage"))


def test_apply_operator_univariate():
    # theta^2 - z applied to log(z) should give -z*log(z), since theta(log z) = 1
    # and theta^2(log z) = theta(1) = 0.
    period = Period(no_variables=1, coefficients={(1,): {(0,): 1}}, order=5)
    op = PFOperator("theta0^2 - z0", no_variables=1)

    result = period.apply_operator(op)

    assert result.coefficients == {(1,): {(1,): -1}}


def test_apply_operator_mixes_variables():
    # z0*theta1 applied to z0*log(z1): theta1 strips log(z1) down to a bare 1
    # (leaving z0 untouched), then multiplying by z0 gives z0^2.
    period = Period(no_variables=2, coefficients={(0, 1): {(1, 0): 1}}, order=5)
    op = PFOperator("z0*theta1", no_variables=2)

    result = period.apply_operator(op)

    assert result.coefficients == {(0, 0): {(2, 0): 1}}


def test_apply_operator_truncates_to_order():
    # Multiplying by z0^2 pushes some terms above the period's order, so they
    # should be dropped rather than kept with a nonzero coefficient.
    period = Period(no_variables=1, coefficients={(0,): {(0,): 1, (1,): 1, (2,): 1}}, order=2)
    op = PFOperator("z0^2", no_variables=1)

    result = period.apply_operator(op)

    assert result.coefficients == {(0,): {(2,): 1}}
    assert result.order == 2


def test_apply_operator_rejects_variable_count_mismatch():
    period = Period(no_variables=1, coefficients={})
    op = PFOperator("z0*theta1", no_variables=2)

    with pytest.raises(ValueError):
        period.apply_operator(op)