# TODO: more
import itertools
from typing import Dict

import numpy as np

from steady_cell_phenotype._util import process_model_text


def test_full_continuity():
    model_text = """
    A=B
    B=A"""

    knockouts: Dict[str, int] = {}  # no knockout
    continuous: Dict[str, bool] = {"A": True, "B": True}  # all variables continuous

    variables, update_fn, equation_system = process_model_text(
        model_text, knockouts, continuous
    )

    assert variables == ("A", "B")

    # small changes unchanged
    for a, b in itertools.product(range(3), repeat=2):
        if abs(a - b) <= 1:
            assert np.all(update_fn([a, b]) == [b, a])

    # all (esp. larger) changes `muted`
    for a, b in itertools.product(range(3), repeat=2):
        assert np.all(
            update_fn([a, b])
            == np.mod(
                [
                    b ** 2 + a + 2 * a ** 2 + a ** 2 * b + 2 * a * b ** 2,
                    a ** 2 + b + 2 * b ** 2 + a * b ** 2 + 2 * a ** 2 * b,
                ],
                3,
            )
        )


def test_partial_continuity():
    model_text = """
    A=B
    B=A"""

    knockouts: Dict[str, int] = {}  # no knockout
    continuous: Dict[str, bool] = {"A": True, "B": False}  # some variables continuous

    variables, update_fn, equation_system = process_model_text(
        model_text, knockouts, continuous
    )

    assert variables == ("A", "B")

    # small changes unchanged
    for a, b in itertools.product(range(3), repeat=2):
        if abs(a - b) <= 1:
            assert np.all(update_fn([a, b]) == [b, a])

    # all (esp. larger) changes half `muted`
    for a, b in itertools.product(range(3), repeat=2):
        assert np.all(
            update_fn([a, b])
            == np.mod([b ** 2 + a + 2 * a ** 2 + a ** 2 * b + 2 * a * b ** 2, a], 3)
        )
