import os

from sage.all import *  # noqa: F401

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


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_power_series_solution_and_recovers_first_operator():
    # Picard--Fuchs ideal for P_{22211}[8] at the MUM point:
    #   D1 = Theta_x^2*(Theta_x - 2*Theta_y) - 4*x*(4*Theta_x + 3)*(4*Theta_x + 2)*(4*Theta_x + 1)
    #   D2 = Theta_y^2 - y*(2*Theta_y - Theta_x + 1)*(2*Theta_y - Theta_x)
    #
    # We test moving the ideal to the intersection with the locus of Strong Coupling.
    D1 = PFOperator(
        "theta0^2*(theta0 - 2*theta1) - 4*z0*(4*theta0 + 3)*(4*theta0 + 2)*(4*theta0 + 1)",
        no_variables=2,
    )
    D2 = PFOperator(
        "theta1^2 - z1*(2*theta1 - theta0 + 1)*(2*theta1 - theta0)",
        no_variables=2,
    )
    moved = PFIdeal([D1, D2]).change_coordinates("z0", "z1 - 1/4")
    op1, op2 = moved.operators

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

    period = solutions[0]
    assert period.apply_operator(op1).coefficients == {}
    assert period.apply_operator(op2).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, 0)] == 0

    recovered = period.find_annihilating_operators(z_degree=1, theta_degree=2)
    assert len(recovered) == 1
    assert recovered[0].operator == PFOperator(
        "-2*z1*theta0^2 + 8*z1*theta0*theta1 - 8*z1*theta1^2 + 2*z1*theta0 - 4*z1*theta1 "
        "+ 2*theta0*theta1 - 2*theta1^2 + theta1",
        no_variables=2,
    ).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