tests/test_er_model.py

"""
Tests for the ER model implementation.
"""
from pathlib import Path

import numpy as np
import pytest
import yaml
import pandas as pd

from er_model_core.allometry import calculate_biomass
from er_model_core.er_model import ERModel, Species
from er_model_core.metrics import calculate_carbon


@pytest.fixture
def sample_config():
    """Create a minimal test configuration."""
    return {
        "species": [
            {
                "name": "Test Species",
                "planting_density": 1000,
                "m_ref": 0.16,
                "DBH_ref": 9.0,
                "p": 1.493,
                "chapman_richards": {
                    "dbh": {"a": 30, "b": 0.15, "c": 1.5},
                    "height": {"a": 10, "b": 0.05, "c": 1.2}
                },
                "allometry": {
                    "biomass_equation": "0.2 * (dbh ** 2.4)",
                    "root_shoot_ratio": 0.4
                },
                "initial_values": {"dbh": 1.0, "height": 0.5}
            }
        ],
        "project": {
            "duration_years": 5,
            "planting_schedule": {
                "year_1": 100
            }
        },
        "carbon": {
            "biomass_to_carbon": 0.47,
            "carbon_to_co2": 3.67,
            "buffer_percentage": 15,
            "leakage_percentage": 0,
            "baseline_emissions": 0
        }
    }


@pytest.fixture
def config_file(tmp_path, sample_config):
    """Create a temporary config file."""
    config_path = tmp_path / "test_config.yaml"
    with open(config_path, "w") as f:
        yaml.dump(sample_config, f)
    return config_path


def test_model_initialization(config_file):
    """Test that the model initializes correctly from config."""
    model = ERModel(config_file)
    assert len(model.species) == 1
    assert model.species[0].name == "Test Species"
    assert model.project.duration_years == 5
    assert model.carbon.biomass_to_carbon == 0.47


def test_surviving_trees_calculation(config_file):
    """Test tree survival calculations with DBH-dependent mortality."""
    model = ERModel(config_file)
    species = model.species[0]
    # Test first year mortality (should use DBH-dependent formula)
    dbh = 9.0
    m = species.m_ref * (species.DBH_ref / dbh) ** species.p
    initial_trees = 1000
    N_live = initial_trees * (1 - m)
    expected = initial_trees * (1 - 0.16)
    assert np.isclose(N_live, expected, atol=1e-3)


def test_chapman_richards_growth(config_file):
    """Test DBH calculations using Chapman-Richards equation."""
    model = ERModel(config_file)
    species = model.species[0]
    
    # Test initial growth
    dbh = model.chapman_richards_dbh(0, species.chapman_richards)
    assert dbh == 0
    
    # Test asymptotic behavior
    dbh = model.chapman_richards_dbh(100, species.chapman_richards)
    assert np.isclose(dbh, species.chapman_richards["a"], rtol=0.01)


def test_biomass_calculation():
    """Test biomass calculations from DBH."""
    params = {
        "biomass_equation": "0.2 * (dbh ** 2.4)",
        "root_shoot_ratio": 0.4
    }
    
    biomass = calculate_biomass(10, params)
    expected_agb = 0.2 * (10 ** 2.4)
    expected_total = expected_agb * (1 + 0.4)
    assert np.isclose(biomass, expected_total)


def test_carbon_calculation():
    """Test carbon conversion calculations."""
    biomass = 1000  # kg
    biomass_to_carbon = 0.47
    carbon_to_co2 = 3.67
    
    co2 = calculate_carbon(biomass, biomass_to_carbon, carbon_to_co2)
    expected = (1000 * 0.47 * 3.67) / 1000  # Convert to metric tons
    assert np.isclose(co2, expected)


def test_full_pipeline(config_file):
    """Test the complete model pipeline."""
    model = ERModel(config_file)
    results = model.run()
    
    assert len(results) == 5  # Duration years
    assert "year" in results.columns
    assert "gross_carbon" in results.columns
    assert "net_carbon" in results.columns
    assert all(results["net_carbon"] <= results["gross_carbon"])


def test_dbh_dependent_mortality():
    """Test that DBH-dependent mortality matches expected at DBH=9 and year-5 plateau."""
    from er_model_core.er_model import ERModel, Species
    # Minimal species config
    species = Species(
        name="Test Species",
        planting_density=1000,
        m_ref=0.16,
        DBH_ref=9.0,
        p=1.493,
        chapman_richards={
            "dbh": {"a": 9.0, "b": 0.5, "c": 1.0},
            "height": {"a": 5.0, "b": 0.2, "c": 1.0}
        },
        allometry={"equation": "0.2 * (dbh ** 2.4)", "root_shoot_ratio": 0.4},
        initial_values={"dbh": 9.0, "height": 5.0}
    )
    # At DBH = 9, mortality = 0.16
    dbh = 9.0
    m = species.m_ref * (species.DBH_ref / dbh) ** species.p
    assert np.isclose(m, 0.16, atol=1e-3)

    # Simulate 5 years, enforce plateau at year 5
    initial_trees = 1000
    plateau_density = 2000
    class DummyModel:
        def chapman_richards_growth(self, age, params, initial):
            return dbh  # Always 9 for this test
    model = DummyModel()
    N_live = initial_trees
    for y in range(1, 6):
        m = species.m_ref * (species.DBH_ref / dbh) ** species.p
        N_live = N_live * (1 - m)
        if y == 5:
            N_live = plateau_density
    assert N_live == plateau_density


def test_allometric_equation_species_A_B():
    """Test Zanvo et al. 2023 allometric equations for both species."""
    from er_model_core.allometry import calculate_biomass
    # Test values
    dbh = 10.0  # cm
    height = 5.0  # m
    # species_A (Rhizophora spp.)
    expected_A = 2.0738 * (dbh**2 * height)**0.67628
    result_A = calculate_biomass(dbh, height, "species_A", {})
    assert np.isclose(result_A, expected_A, rtol=1e-6), f"species_A: got {result_A}, expected {expected_A}"
    # species_B (Avicennia germinans)
    expected_B = 1.5595 * (dbh**2 * height)**0.55864
    result_B = calculate_biomass(dbh, height, "species_B", {})
    assert np.isclose(result_B, expected_B, rtol=1e-6), f"species_B: got {result_B}, expected {expected_B}" 

# --- Parameter sweep/test for plausible survival curves (moved from er_model.py) ---
def test_dbh_mortality_sweep():
    import matplotlib.pyplot as plt
    import numpy as np
    m_refs = [0.01, 0.05, 0.1, 0.16]
    ps = [1.0, 1.5, 2.0]
    DBH_ref = 9.0
    years = np.arange(1, 31)
    initial_trees = 1000
    results = {}
    for m_ref in m_refs:
        for p in ps:
            N_live = initial_trees
            N_lives = []
            for year in years:
                # This test requires a simplified DBH growth for its internal logic.
                # It does not use the main ERModel's growth functions directly.
                dbh_test_growth = 1.0 + (year - 1) * 0.5  # simple linear DBH growth for test
                dbh = max(dbh_test_growth, 1.0)
                
                # Mortality calculation as per ERModel logic being tested
                m = m_ref * (DBH_ref / dbh) ** p
                m = min(max(m, 0), 0.99) 
                N_live = N_live * (1 - m)
                N_lives.append(N_live)
            results[(m_ref, p)] = N_lives

    # This test originally had plotting. For automated tests, assertions are better.
    # If visual inspection is needed, this part can be run in a notebook.
    # For now, let's assert that the number of trees is non-increasing.
    for (m_ref, p), N_lives_curve in results.items():
        for i in range(len(N_lives_curve) - 1):
            assert N_lives_curve[i+1] <= N_lives_curve[i], f"Mortality sweep {m_ref, p}: Trees increased from {N_lives_curve[i]} to {N_lives_curve[i+1]}"
        assert N_lives_curve[-1] < initial_trees, f"Mortality sweep {m_ref, p}: No mortality occurred or trees increased."

    # plt.figure(figsize=(10,6))
    # for (m_ref, p), N_lives_plot in results.items():
    #     plt.plot(years, N_lives_plot, label=f"m_ref={m_ref}, p={p}")
    # plt.xlabel("Year")
    # plt.ylabel("Surviving Trees")
    # plt.title("DBH-dependent Mortality Parameter Sweep")
    # plt.legend()
    # plt.grid(True)
    # To save plot instead of showing:
    # output_dir = Path(__file__).parent.parent / "outputs" / "test_plots"
    # output_dir.mkdir(parents=True, exist_ok=True)
    # plt.savefig(output_dir / "dbh_mortality_sweep.png")
    # plt.close()
    print("test_dbh_mortality_sweep completed basic assertions.")

# If you want to run this specific test and see the plot (e.g., during development):
# if __name__ == "__main__":
#     test_dbh_mortality_sweep() # You would need to handle imports for ERModel parts if called directly here