"""Stage `fields` and shared zone modifiers."""
from __future__ import annotations

import numpy as np

from .config import MASK_DEFAULTS, params


def gravity_mod(cfg, m_gravity):
    M = params(cfg, "masks", MASK_DEFAULTS)
    return np.clip(1.0 + M["gravity_range"] * np.asarray(m_gravity, dtype=np.float64), 0.05, None)


def o2_mod(cfg, m_o2):
    M = params(cfg, "masks", MASK_DEFAULTS)
    return np.clip(1.0 + M["o2_range"] * np.asarray(m_o2, dtype=np.float64), 0.0, None)


def pressure(p0_bar, z_m, scale_height_m):
    """Air pressure at ground height z (bar): sea-level pressure p0 falling with height (below 0 m: p0)."""
    return p0_bar * np.exp(-np.maximum(np.asarray(z_m, dtype=np.float64), 0.0) / scale_height_m)


def plant_height(g_planet, gravity_g):
    """How much taller plants grow than under the planet's normal gravity (height limit ∝ 1 / g)."""
    return g_planet / np.asarray(gravity_g, dtype=np.float64)


def run(ctx) -> dict:
    pl = ctx.cfg["planet"]
    zs, m_o2, m_g = ctx.need("z_surface_m", "m_o2_zones", "m_gravity_zones")
    p = pressure(pl["sea_level_pressure_bar"], zs, pl["scale_height_km"] * 1000.0)
    o2 = pl["o2_fraction"] * o2_mod(ctx.cfg, m_o2)
    return {"po2_bar": o2 * p, "gravity_g": pl["gravity_g"] * gravity_mod(ctx.cfg, m_g), "pressure_bar": p,
            "o2_fraction": o2, "fire_reactivity": np.ones(len(p)),   # zone events (eras) set other values
            "plant_height_x": plant_height(pl["gravity_g"], pl["gravity_g"] * gravity_mod(ctx.cfg, m_g))}