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_find_annihilating_operators_recovers_theta_squared():
    # theta0^2 annihilates log(z0): theta0(log z0) = 1, theta0^2(log z0) = theta0(1) = 0.
    # Among degree-(0, 2) Ansatze this should be the only solution, up to scaling.
    period = Period(no_variables=1, coefficients={(1,): {(0,): 1}}, order=5)

    ops = period.find_annihilating_operators(z_degree=0, theta_degree=2)

    # A one-dimensional solution space means exactly one basis operator.
    assert len(ops) == 1
    assert period.apply_operator(ops[0]).coefficients == {}
    assert str(ops[0].operator) == "theta0^2"


def test_find_annihilating_operators_recovers_geometric_series_operator():
    # sum_{k=0}^{4} z0^k is annihilated (up to truncation order) by
    # (1 - z0)*theta0 - z0, i.e. theta0 - z0*theta0 - z0.
    period = Period(
        no_variables=1,
        coefficients={(0,): {(0,): 1, (1,): 1, (2,): 1, (3,): 1, (4,): 1}},
        order=4,
    )

    ops = period.find_annihilating_operators(z_degree=1, theta_degree=1)

    assert len(ops) == 1
    assert period.apply_operator(ops[0]).coefficients == {}


def test_find_annihilating_operators_returns_empty_list_when_no_solution_exists():
    # No degree-0 (constant) operator other than the zero operator can annihilate a
    # nonzero constant period, so the linear system's only solution is trivial.
    period = Period(no_variables=1, coefficients={(0,): {(0,): 1}}, order=0)

    ops = period.find_annihilating_operators(z_degree=0, theta_degree=0)

    assert ops == []


def test_simplify_factorises_theta_polynomial_per_z_monomial():
    # theta0^4 - 5*z0*(5*theta0+1)*(5*theta0+2)*(5*theta0+3)*(5*theta0+4), expanded, is the
    # quintic's Picard-Fuchs operator. simplify() should recover the factorised form: the
    # z0^0 part (theta0^4) has no theta-factor to pull out, while the z0^1 part factorises
    # into the four linear pieces.
    expanded = "theta0^4 - 3125*z0*theta0^4 - 6250*z0*theta0^3 - 4375*z0*theta0^2 - 1250*z0*theta0 - 120*z0"
    op = PFOperator(expanded, no_variables=1)

    simplified = op.simplify()

    # The underlying (expanded) operator is unchanged - only the display string differs.
    assert simplified.operator == op.operator
    assert simplified.operator_string == "-5*(5*theta0 + 4)*(5*theta0 + 3)*(5*theta0 + 2)*(5*theta0 + 1)*z0 + theta0^4"


def test_find_power_series_solution_recovers_quintic_period():
    # The quintic's Picard-Fuchs operator has indicial equation theta0^4 = 0 at z0 = 0
    # (a quadruple root at 0), so its holomorphic (non-logarithmic) power series solution
    # is found at indicial 0. Up to normalisation this is the classic quintic period
    # 1 + 120*z0 + 113400*z0^2 + ...
    op = PFOperator(
        "theta0^4 - 5*z0*(5*theta0 + 1)*(5*theta0 + 2)*(5*theta0 + 3)*(5*theta0 + 4)",
        no_variables=1,
    )
    ideal = PFIdeal([op])

    solutions = ideal.find_power_series_solution(indicials=[0], order=3)

    assert len(solutions) == 1
    assert solutions[0].coefficients == {(0,): {(0,): 1, (1,): 120, (2,): 113400, (3,): 168168000}}


def test_find_power_series_solution_handles_rational_indicial():
    # theta0 - 1/2 kills z0^(1/2 + k) only when k = 0, since its eigenvalue is 1/2 + k;
    # so at indicial 1/2 there is exactly one solution (a constant multiple of sqrt(z0)),
    # while at indicial 0 no power series solution exists at all.
    op = PFOperator("theta0 - 1/2", no_variables=1)
    ideal = PFIdeal([op])

    solutions = ideal.find_power_series_solution(indicials=[1/2], order=3)
    assert len(solutions) == 1
    assert solutions[0].coefficients == {(0,): {(0,): 1, (1,): 0, (2,): 0, (3,): 0}}
    assert solutions[0].period_string == "sqrt(z0)"

    assert ideal.find_power_series_solution(indicials=[0], order=3) == []


def test_ideal_finds_holomorphic_solution_and_recovers_first_operator():
    # Picard--Fuchs ideal for P_{22211}[8]:
    #   M1 = theta2*(-2*theta1+2*theta2-1) + 2*(theta1-2*theta2-1)*(theta1-2*theta2)*z2
    #   M2 = theta1^2*(2*(theta1-2*theta2)*z2-theta2) - 16*(2*theta1+1)*(4*theta1+1)*(4*theta1+3)*z1*z2
    M1 = PFOperator(
        "theta1*(-2*theta0 + 2*theta1 - 1) + 2*(theta0 - 2*theta1 - 1)*(theta0 - 2*theta1)*z1",
        no_variables=2,
    )
    M2 = PFOperator(
        "theta0^2*(2*(theta0 - 2*theta1)*z1 - theta1) - 16*(2*theta0 + 1)*(4*theta0 + 1)*(4*theta0 + 3)*z0*z1",
        no_variables=2,
    )
    ideal = PFIdeal([M1, M2])

    solutions = ideal.find_power_series_solution(indicials=[0, QQ(1) / 2], order=6)
    assert len(solutions) == 1

    period = solutions[0]
    assert period.apply_operator(M1).coefficients == {}
    assert period.apply_operator(M2).coefficients == {}
    # Normalisation convention: the free parameter at the leading (0, 0) coefficient is 1.
    assert period.coefficients[(0, 0)][(0, 0)] == 1
    assert period.coefficients[(0, 0)][(1, 1)] == -32

    # M1 lives entirely within z_degree <= 1, theta_degree <= 2 (a single bare factor of
    # z2, quadratic in theta), so this is the natural Ansatz level to look for it at.
    recovered = period.find_annihilating_operators(z_degree=1, theta_degree=2)
    assert len(recovered) == 1

    # recovered[0] should be a scalar multiple of M1 - compare via the coefficient of the
    # bare theta2 (theta1 in code) monomial, which is nonzero in M1.
    ratio = (
        recovered[0].operator.monomial_coefficient(M1.theta_gens[1])
        / M1.operator.monomial_coefficient(M1.theta_gens[1])
    )
    assert ratio != 0
    assert recovered[0].operator == ratio * M1.operator


def test_single_operator_recovers_itself_from_its_holomorphic_period():
    # Ensuring methods work for operator 4.2.1:
    L = PFOperator(
        "theta0^4 - 4*z0*(2*theta0 + 1)^2*(7*theta0^2 + 7*theta0 + 2) "
        "- 128*z0^2*(2*theta0 + 1)^2*(2*theta0 + 3)^2",
        no_variables=1,
    )
    ideal = PFIdeal([L])

    solutions = ideal.find_power_series_solution(indicials=[0], order=15)
    assert len(solutions) == 1

    period = solutions[0]
    assert period.apply_operator(L).coefficients == {}
    assert period.coefficients[(0,)][(0,)] == 1
    assert period.coefficients[(0,)][(1,)] == 8
    assert period.coefficients[(0,)][(2,)] == 360

    # L lives entirely within z_degree <= 2, theta_degree <= 4, its stated level.
    recovered = period.find_annihilating_operators(z_degree=2, theta_degree=4)
    assert len(recovered) == 1

    # recovered[0] should be a scalar multiple of L - compare via the coefficient of the
    # bare theta0^4 monomial, which is 1 in L.
    ratio = recovered[0].operator.monomial_coefficient(L.theta_gens[0] ** 4) / L.operator.monomial_coefficient(
        L.theta_gens[0] ** 4
    )
    assert ratio != 0
    assert recovered[0].operator == ratio * L.operator