"""
Run a Model
============================

.. sectionauthor:: Julian Heinze (Helmholtz Centre for Environmental Research GmbH - UFZ)

This example demonstrates basic usage of the Spatiocoexistence model.
"""

# %%
# Necessary imports
from spatiocoexistence import Model, Inventory
from spatiocoexistence.data import inventory_file, parameter_file

# %%
# Initialize the model with inventory and parameter files
# This model will be simulated on a single thread.
# If an integer greater than 1 is used for 'threads', the simulation will run in parallel on that many threads.
inventory = Inventory.from_data(inventory_file)
model = Model(initial_inventory=inventory, parameters=parameter_file, threads=1)

# %%
# Get the current parameter values and modify them
parameter_old = model.parameters
# update the minimal survival probability of the recruits, increase the radius to 20m and the dispersal range for seeds to 25m
model.update_parameters(
    {"min_recruit_survival": 0.9, "neighborhood_radius": 20, "sigma_dispersal": 25.0}
)

print(
    "Original parameters (subset):",
    {
        k: parameter_old[k]
        for k in ["neighborhood_radius", "sigma_dispersal", "min_recruit_survival"]
    },
)
print(
    "Updated parameters (subset):",
    {
        k: model.parameters[k]
        for k in ["neighborhood_radius", "sigma_dispersal", "min_recruit_survival"]
    },
)

# %%
# Run a short simulation of 100 steps (each step represents 5 years)
model.step(100)

# %%
# Access the updated inventory after simulation
inv = model.inventory

# Basic analyses on the resulting inventory
num_individuals = inv.x.shape[0]
mean_dbh = inv.dbh.mean()
max_species_id = inv.species.max()
abundances = inv.species_abundance()

print(f"Number of individuals after 100 steps: {num_individuals}")
print(f"Mean DBH (dm) after 100 steps: {mean_dbh:.3f}")
print("Species abundance distribution (counts per species):", abundances[0])


# %%
# Show plot of the current inventory state.
inv.plot(threshold=50, filter=1.0)

# %%
# Demonstrate individual processes on the current state
# Store initial inventory size for comparison
initial_size = len(model.inventory.data)
print(f"Initial inventory size: {initial_size}")

# %%
# Apply one growth update and show mean dbh change
prev_mean_dbh = model.inventory.dbh.mean()
delta_dbh = model.growth()
post_growth_mean_dbh = model.inventory.dbh.mean()
print(
    f"Mean DBH before growth: {prev_mean_dbh:.3f} -> after growth: {post_growth_mean_dbh:.3f}"
)
print(f"Mean DBH increase: {delta_dbh.mean():.4f}")

# %%
# Apply mortality and report how many died
# Note: Dead individuals are now automatically removed from the inventory
deads_inventory = model.mortality()
print(
    f"Mortality event: {deads_inventory.data.shape[0]} individuals died and were removed."
)
print(f"Inventory size after mortality: {model.inventory.data.shape[0]}")
deads_inventory.plot_map(filter=0)
# %%
# Apply recruitment and get new recruits as an Inventory object
# The recruits are automatically added to the main inventory
recruits_inventory = model.recruitment()
print(f"\nRecruitment event: {len(recruits_inventory.data)} new individuals recruited")
print(f"Inventory size after recruitment: {len(model.inventory.data)}")
recruits_inventory.plot_map(filter=0, scale=10)