"""Test fractional product handling in Mobspy."""

import pytest
import mobspy
from mobspy import (
    BaseSpecies,
    Simulation,
    New,
    u,
    ModelParameters,
    simlog,
    Assign,
    Rev,
    All,
    set_counts,
)


def test_deterministic():
    """Test fractional products in deterministic simulation."""
    A, B, C = BaseSpecies()
    A(10)
    B(1)
    A + B >> 0.5 * C + B [0.005]
    Sim = Simulation(A | B | C)
    Sim.method = "deterministic"
    Sim.volume = 1 * u.mL
    Sim.run(duration=10, unit_x=u.s, unit_y=1 / u.mL, plot_data=False)
    
    # approximate check
    assert Sim.results["C"][0][-1] == pytest.approx(5)

def test_deterministic_variable():
    """Test fractional products in deterministic simulation.
    The yield coefficient is defined as a variable."""
    A, B, C = BaseSpecies()
    A(10)
    B(1)
    yield_coeff = 0.5
    A + B >> yield_coeff * C + B [0.005]
    Sim = Simulation(A | B | C)
    Sim.method = "deterministic"
    Sim.volume = 1 * u.mL
    Sim.run(duration=10, unit_x=u.s, unit_y=1 / u.mL, plot_data=False)
    
    # approximate check
    assert Sim.results["C"][0][-1] == pytest.approx(5)


def test_stochastic():
    """Test fractional products in stochastic simulation."""
    A, B, C = BaseSpecies()
    A(10)
    B(1)
    A + B >> 0.5 * C + B [0.005]
    Sim = Simulation(A | B | C)
    Sim.method = "stochastic"
    Sim.volume = 1 * u.mL
    Sim.run(duration=10, unit_x=u.s, unit_y=1 / u.mL, plot_data=False)
    
    # approximate check
    assert Sim.results["C"][0][-1] == pytest.approx(5)

def test_stochastic_variable():
    """Test fractional products in stochastic simulation.
    The yield coefficient is defined as a variable."""
    # TODO: currently coefficients 0.5 to 1.0 + integer are supported
    A, B, C = BaseSpecies()
    A(10)
    B(1)
    yield_coeff = 0.5
    A + B >> yield_coeff * C + B [0.005]
    Sim = Simulation(A | B | C)
    Sim.method = "stochastic"
    Sim.volume = 1 * u.mL
    Sim.run(duration=10, unit_x=u.s, unit_y=1 / u.mL, plot_data=False)
    
    # approximate check
    assert Sim.results["C"][0][-1] == pytest.approx(5)


if __name__ == "__main__":
    test_deterministic()
    test_stochastic()
    test_deterministic_variable()
    test_stochastic_variable()
