Zalmoxis.zalmoxis

def parse_eos_config(eos_section):
    """Parse [EOS] TOML section into per-layer EOS dict.

    Info
    -----
    Supports new per-layer format (core/mantle/ice_layer fields) and legacy
    format (choice field) for backward compatibility.

    Parameters
    ----------
    eos_section : dict
        The [EOS] section from the TOML config.

    Returns
    -------
    dict
        Per-layer EOS config, e.g.
        {"core": "Seager2007:iron", "mantle": "WolfBower2018:MgSiO3"}.
    """
    # New format: per-layer fields present
    if 'core' in eos_section:
        if 'mantle' not in eos_section:
            raise ValueError(
                "EOS config has 'core' but missing 'mantle'. "
                "Both 'core' and 'mantle' are required."
            )
        layer_eos = {
            'core': eos_section['core'],
            'mantle': eos_section['mantle'],
        }
        water = eos_section.get('ice_layer', '')
        if water:
            layer_eos['ice_layer'] = water
        return layer_eos

    # Legacy format: expand 'choice' field
    choice = eos_section.get('choice', '')
    if choice in LEGACY_EOS_MAP:
        return dict(LEGACY_EOS_MAP[choice])

    if choice == 'Analytic:Seager2007':
        layer_eos = {
            'core': f'Analytic:{eos_section.get("core_material", "iron")}',
            'mantle': f'Analytic:{eos_section.get("mantle_material", "MgSiO3")}',
        }
        water_mat = eos_section.get('water_layer_material', '')
        if water_mat:
            layer_eos['ice_layer'] = f'Analytic:{water_mat}'
        return layer_eos

    raise ValueError(
        f"Unknown EOS config. Set per-layer fields (core, mantle) or legacy 'choice'. "
        f'Got: {eos_section}'
    )

def validate_layer_eos(layer_eos_config):
    """Validate all per-layer EOS strings, including multi-material mixtures.

    Parameters
    ----------
    layer_eos_config : dict
        Per-layer EOS config from parse_eos_config().

    Raises
    ------
    ValueError
        If any layer EOS string or component is invalid.
    """
    for layer, eos_str in layer_eos_config.items():
        mixture = parse_layer_components(eos_str)
        for comp in mixture.components:
            if comp in VALID_TABULATED_EOS:
                continue
            if comp.startswith('Analytic:'):
                material_key = comp.split(':', 1)[1]
                if material_key not in VALID_MATERIAL_KEYS:
                    raise ValueError(
                        f"Invalid analytic material '{material_key}' "
                        f"in layer '{layer}'. "
                        f'Valid keys: {sorted(VALID_MATERIAL_KEYS)}'
                    )
                continue
            raise ValueError(
                f"Invalid EOS component '{comp}' in layer '{layer}'. "
                f'Valid tabulated: {sorted(VALID_TABULATED_EOS)}. '
                f'Valid analytic: Analytic:<material> with '
                f'{sorted(VALID_MATERIAL_KEYS)}.'
            )

def choose_config_file(temp_config_path=None):
    """
    Function to choose the configuration file to run the main function.

    Info
    -----
    The function will first check if a temporary configuration file is provided.
    If not, it will check if the -c flag is provided in the command line arguments.
    If the -c flag is provided, the function will read the configuration file path from the next argument.
    If no temporary configuration file or -c flag is provided, the function will read the default configuration file.

    Parameters
    ----------
    temp_config_path : str, optional
        Path to a temporary configuration file. If provided, this file will be used instead of the default or -c specified config.

    Returns
    -------
    dict
        The loaded configuration parameters from the chosen config file.
    """

    # Load the configuration file either from terminal (-c flag) or default path
    if temp_config_path:
        try:
            config = toml.load(temp_config_path)
            logger.info(f'Reading temporary config file from: {temp_config_path}')
        except FileNotFoundError:
            logger.error(f'Error: Temporary config file not found at {temp_config_path}')
            sys.exit(1)
    elif '-c' in sys.argv:
        index = sys.argv.index('-c')
        try:
            config_file_path = sys.argv[index + 1]
            config = toml.load(config_file_path)
            logger.info(f'Reading config file from: {config_file_path}')
        except IndexError:
            logger.error('Error: -c flag provided but no config file path specified.')
            sys.exit(1)  # Exit with error code
        except FileNotFoundError:
            logger.error(f'Error: Config file not found at {config_file_path}')
            sys.exit(1)
    else:
        config_default_path = os.path.join(ZALMOXIS_ROOT, 'input', 'default.toml')
        try:
            config = toml.load(config_default_path)
            logger.info(f'Reading default config file from {config_default_path}')
        except FileNotFoundError:
            logger.info(f'Error: Default config file not found at {config_default_path}')
            sys.exit(1)

    return config

def load_zalmoxis_config(temp_config_path=None):
    """Load and return configuration parameters for the Zalmoxis model.

    Returns
    -------
    dict
        All relevant configuration parameters.
    """
    config = choose_config_file(temp_config_path)

    # Parse per-layer EOS config (supports both new and legacy formats)
    layer_eos_config = parse_eos_config(config['EOS'])
    validate_layer_eos(layer_eos_config)

    # Melting curve config (defaults for backward compat with old TOML files)
    eos_section = config['EOS']
    rock_solidus = eos_section.get('rock_solidus', 'Stixrude14-solidus')
    rock_liquidus = eos_section.get('rock_liquidus', 'Stixrude14-liquidus')
    mushy_zone_factor = eos_section.get('mushy_zone_factor', 1.0)
    condensed_rho_min = eos_section.get('condensed_rho_min', CONDENSED_RHO_MIN_DEFAULT)
    condensed_rho_scale = eos_section.get('condensed_rho_scale', CONDENSED_RHO_SCALE_DEFAULT)
    binodal_T_scale = eos_section.get('binodal_T_scale', BINODAL_T_SCALE_DEFAULT)

    # Build per-EOS mushy_zone_factors dict. Only include materials that are
    # actually configured in a layer; unused materials default to 1.0 so that
    # a global mushy_zone_factor < 1.0 does not trigger the validation check
    # for materials absent from the model.
    _paleos_materials = {
        'PALEOS:iron': 'mushy_zone_factor_iron',
        'PALEOS:MgSiO3': 'mushy_zone_factor_MgSiO3',
        'PALEOS:H2O': 'mushy_zone_factor_H2O',
    }
    # Collect all EOS component strings from all layers
    _all_eos_strings = ' '.join(v for v in layer_eos_config.values() if v)
    mushy_zone_factors = {}
    for paleos_name, toml_key in _paleos_materials.items():
        if paleos_name in _all_eos_strings:
            # Material is in use: apply per-material override or global default
            mushy_zone_factors[paleos_name] = eos_section.get(toml_key, mushy_zone_factor)
        else:
            # Material not in use: default to 1.0 (no mushy zone)
            mushy_zone_factors[paleos_name] = eos_section.get(toml_key, 1.0)

    config_params = {
        'planet_mass': config['InputParameter']['planet_mass'] * earth_mass,
        'core_mass_fraction': config['AssumptionsAndInitialGuesses']['core_mass_fraction'],
        'mantle_mass_fraction': config['AssumptionsAndInitialGuesses']['mantle_mass_fraction'],
        'temperature_mode': config['AssumptionsAndInitialGuesses']['temperature_mode'],
        'surface_temperature': config['AssumptionsAndInitialGuesses']['surface_temperature'],
        'center_temperature': config['AssumptionsAndInitialGuesses']['center_temperature'],
        'temp_profile_file': config['AssumptionsAndInitialGuesses']['temperature_profile_file'],
        'layer_eos_config': layer_eos_config,
        'rock_solidus': rock_solidus,
        'rock_liquidus': rock_liquidus,
        'mushy_zone_factor': mushy_zone_factor,
        'mushy_zone_factors': mushy_zone_factors,
        'condensed_rho_min': condensed_rho_min,
        'condensed_rho_scale': condensed_rho_scale,
        'binodal_T_scale': binodal_T_scale,
        'num_layers': config['Calculations']['num_layers'],
        'max_iterations_outer': config['IterativeProcess']['max_iterations_outer'],
        'tolerance_outer': config['IterativeProcess']['tolerance_outer'],
        'max_iterations_inner': config['IterativeProcess']['max_iterations_inner'],
        'tolerance_inner': config['IterativeProcess']['tolerance_inner'],
        'relative_tolerance': config['IterativeProcess']['relative_tolerance'],
        'absolute_tolerance': config['IterativeProcess']['absolute_tolerance'],
        'maximum_step': config['IterativeProcess']['maximum_step'],
        'adaptive_radial_fraction': config['IterativeProcess']['adaptive_radial_fraction'],
        'max_center_pressure_guess': config['IterativeProcess']['max_center_pressure_guess'],
        'target_surface_pressure': config['PressureAdjustment']['target_surface_pressure'],
        'pressure_tolerance': config['PressureAdjustment']['pressure_tolerance'],
        'max_iterations_pressure': config['PressureAdjustment']['max_iterations_pressure'],
        'data_output_enabled': config['Output']['data_enabled'],
        'plotting_enabled': config['Output']['plots_enabled'],
        'verbose': config['Output']['verbose'],
        'iteration_profiles_enabled': config['Output']['iteration_profiles_enabled'],
    }

    # Validate all parameters for physical and logical consistency
    validate_config(config_params)

    return config_params

def load_material_dictionaries():
    """Load and return the material properties dictionaries.

    Returns
    -------
    dict
        EOS registry keyed by EOS identifier string (e.g.
        ``"Seager2007:iron"``, ``"PALEOS:MgSiO3"``).
    """
    return dict(EOS_REGISTRY)

def load_solidus_liquidus_functions(
    layer_eos_config,
    solidus_id='Stixrude14-solidus',
    liquidus_id='Stixrude14-liquidus',
):
    """Load solidus and liquidus functions if any layer uses an EOS that needs them.

    Unified PALEOS tables (PALEOS:iron, PALEOS:MgSiO3, PALEOS:H2O) are
    T-dependent but derive their phase boundary from the table itself, so
    they do not need external melting curves.

    Parameters
    ----------
    layer_eos_config : dict
        Per-layer EOS config.
    solidus_id : str
        Solidus melting curve identifier.
    liquidus_id : str
        Liquidus melting curve identifier.

    Returns
    -------
    tuple or None
        (solidus_func, liquidus_func) if needed, else None.
    """
    all_comps = set()
    for v in layer_eos_config.values():
        if v:
            m = parse_layer_components(v)
            all_comps.update(m.components)
    if all_comps & _NEEDS_MELTING_CURVES:
        solidus_func, liquidus_func = get_solidus_liquidus_functions(solidus_id, liquidus_id)
        return (solidus_func, liquidus_func)
    return None

def main(
    config_params,
    material_dictionaries,
    melting_curves_functions,
    input_dir,
    layer_mixtures=None,
):
    """Run the exoplanet internal structure model.

    Parameters
    ----------
    config_params : dict
        Configuration parameters for the model.
    material_dictionaries : dict
        EOS registry dict keyed by EOS identifier string.
    melting_curves_functions : tuple or None
        (solidus_func, liquidus_func) for EOS needing external melting curves.
    input_dir : str
        Directory containing input files.
    layer_mixtures : dict or None, optional
        Per-layer LayerMixture objects. If None, parsed from
        ``config_params['layer_eos_config']``. PROTEUS/CALLIOPE can provide
        pre-built mixtures with runtime-updated fractions.

    Returns
    -------
    dict
        Model results including radii, density, gravity, pressure, temperature,
        mass enclosed, convergence status, and timing.
    """
    # Initialize convergence flags
    converged = False
    converged_pressure = False
    converged_density = False
    converged_mass = False

    # Unpack configuration parameters
    planet_mass = config_params['planet_mass']
    core_mass_fraction = config_params['core_mass_fraction']
    mantle_mass_fraction = config_params['mantle_mass_fraction']
    temperature_mode = config_params['temperature_mode']
    surface_temperature = config_params['surface_temperature']
    center_temperature = config_params['center_temperature']
    temp_profile_file = config_params['temp_profile_file']
    layer_eos_config = config_params['layer_eos_config']
    num_layers = config_params['num_layers']
    max_iterations_outer = config_params['max_iterations_outer']
    tolerance_outer = config_params['tolerance_outer']
    max_iterations_inner = config_params['max_iterations_inner']
    tolerance_inner = config_params['tolerance_inner']
    relative_tolerance = config_params['relative_tolerance']
    absolute_tolerance = config_params['absolute_tolerance']
    maximum_step = config_params['maximum_step']
    adaptive_radial_fraction = config_params['adaptive_radial_fraction']
    max_center_pressure_guess = config_params['max_center_pressure_guess']
    target_surface_pressure = config_params['target_surface_pressure']
    pressure_tolerance = config_params['pressure_tolerance']
    max_iterations_pressure = config_params['max_iterations_pressure']
    verbose = config_params['verbose']
    iteration_profiles_enabled = config_params['iteration_profiles_enabled']
    # Build per-EOS mushy_zone_factors dict. Prefer the dict if present
    # (set by load_zalmoxis_config). Fall back to building one from the
    # single float for backward compat with callers that only set the
    # global 'mushy_zone_factor' key.
    if 'mushy_zone_factors' in config_params:
        mushy_zone_factors = config_params['mushy_zone_factors']
    else:
        _global_mzf = config_params.get('mushy_zone_factor', 1.0)
        mushy_zone_factors = {
            'PALEOS:iron': _global_mzf,
            'PALEOS:MgSiO3': _global_mzf,
            'PALEOS:H2O': _global_mzf,
        }
    condensed_rho_min = config_params.get('condensed_rho_min', CONDENSED_RHO_MIN_DEFAULT)
    condensed_rho_scale = config_params.get('condensed_rho_scale', CONDENSED_RHO_SCALE_DEFAULT)
    binodal_T_scale = config_params.get('binodal_T_scale', BINODAL_T_SCALE_DEFAULT)

    # Parse layer mixtures if not provided externally (PROTEUS/CALLIOPE)
    if layer_mixtures is None:
        layer_mixtures = parse_all_layer_mixtures(layer_eos_config)

    # Check if any component in any layer uses T-dependent EOS
    uses_Tdep = any_component_is_tdep(layer_mixtures)

    # Enforce per-EOS mass limits for T-dependent tables
    mass_in_earth = planet_mass / earth_mass
    all_comps = set()
    for mix in layer_mixtures.values():
        all_comps.update(mix.components)
    tdep_comps = all_comps & TDEP_EOS_NAMES
    if tdep_comps:
        for eos_name in tdep_comps:
            if eos_name == 'WolfBower2018:MgSiO3':
                max_mass = WOLFBOWER2018_MAX_MASS_EARTH
                reason = (
                    'Deep-mantle pressures far exceed the 1 TPa table boundary '
                    'at higher masses, making clamped densities unreliable. '
                    'Use RTPress100TPa:MgSiO3 for higher masses.'
                )
            elif eos_name == 'RTPress100TPa:MgSiO3':
                max_mass = RTPRESS100TPA_MAX_MASS_EARTH
                reason = (
                    'The RTPress100TPa melt table extends to 100 TPa but '
                    'the solid table is limited to 1 TPa.'
                )
            elif eos_name == 'PALEOS-2phase:MgSiO3':
                max_mass = PALEOS_MAX_MASS_EARTH
                reason = (
                    'The PALEOS MgSiO3 tables extend to 100 TPa for both '
                    'solid and liquid phases.'
                )
            elif eos_name in ('PALEOS:iron', 'PALEOS:MgSiO3', 'PALEOS:H2O'):
                max_mass = PALEOS_UNIFIED_MAX_MASS_EARTH
                reason = 'The unified PALEOS tables extend to 100 TPa (P: 1 bar to 100 TPa).'
            elif eos_name == 'Chabrier:H':
                max_mass = PALEOS_UNIFIED_MAX_MASS_EARTH
                reason = (
                    'The Chabrier H table extends to 10^22 Pa but '
                    'has only been validated up to ~50 M_earth.'
                )
            else:
                continue
            if mass_in_earth > max_mass:
                raise ValueError(
                    f'{eos_name} EOS is limited to planets <= '
                    f'{max_mass} M_earth (requested {mass_in_earth:.2f} M_earth). '
                    f'{reason}'
                )

    # Setup initial guesses
    radius_guess = (
        1000 * (7030 - 1840 * core_mass_fraction) * (planet_mass / earth_mass) ** 0.282
    )
    cmb_mass = 0
    core_mantle_mass = 0

    logger.info(
        f'Starting structure model for a {planet_mass / earth_mass} Earth masses planet '
        f"with EOS config {layer_eos_config} and temperature mode '{temperature_mode}'."
    )

    start_time = time.time()

    # Initialize empty cache for interpolation functions
    interpolation_cache = {}

    # Load solidus and liquidus functions if the caller provided them.
    # Unified PALEOS tables don't need external melting curves, so
    # melting_curves_functions may be None even when uses_Tdep is True.
    if melting_curves_functions is not None:
        solidus_func, liquidus_func = melting_curves_functions
    else:
        solidus_func, liquidus_func = None, None

    # --- Adiabatic temperature mode (standalone Zalmoxis) ----------------
    #
    # When temperature_mode='adiabatic', Zalmoxis computes a self-consistent
    # T(r) from EOS adiabat gradient tables. The transition from the initial
    # linear-T guess to the full adiabat is GRADUAL via a blending parameter:
    #
    #   blend = 0.0   (iteration 0: pure linear T)
    #   blend = 0.5   (first post-convergence iteration: half adiabat)
    #   blend = 1.0   (second post-convergence iteration: full adiabat)
    #
    # This blending prevents the solver from diverging when the temperature
    # profile changes abruptly from linear to adiabatic.
    #
    # In the PROTEUS-SPIDER coupling, temperature_mode is typically
    # 'adiabatic' in the config, but the blend never activates because
    # the initial linear-T structure converges and the mass break fires
    # with blend=0. This is correct: SPIDER provides its own T(r).
    # -------------------------------------------------------------------

    # Storage for the previous iteration's converged profiles.
    # Used by adiabatic mode to compute T(r) from the last P(r) and M(r).
    prev_radii = None
    prev_pressure = None
    prev_mass_enclosed = None

    # Adiabat blending state.
    _using_adiabat = False
    _adiabat_blend = 0.0
    _ADIABAT_BLEND_STEP = 0.25

    # For adiabatic mode, cap the initial center_temperature guess to
    # prevent the linear-T initial profile from being far above the
    # actual adiabat. A typical rocky planet adiabat at 1 M_earth has
    # T_center ~ 3 * T_surface. If center_temperature >> this, the
    # blend from linear to adiabat causes a huge density perturbation
    # that destabilizes convergence. This is conservative (adiabats
    # can be steeper for massive planets), so we use 5 * T_surface.
    if temperature_mode == 'adiabatic':
        max_reasonable_T_center = max(5.0 * surface_temperature, 3000.0)
        if center_temperature > max_reasonable_T_center:
            center_temperature = max_reasonable_T_center
            verbose and logger.info(
                f'Adiabatic mode: capped center_temperature initial guess '
                f'to {center_temperature:.0f} K (5x surface or 3000 K).'
            )

    # Solve the interior structure
    for outer_iter in range(max_iterations_outer):
        radii = np.linspace(0, radius_guess, num_layers)

        density = np.zeros(num_layers)
        mass_enclosed = np.zeros(num_layers)
        gravity = np.zeros(num_layers)
        pressure = np.zeros(num_layers)

        if uses_Tdep:
            # Compute the linear (initial guess) temperature profile.
            # This is a function of radius only; wrap it to accept (r, P).
            _linear_tf = calculate_temperature_profile(
                radii,
                'linear',
                surface_temperature,
                center_temperature,
                input_dir,
                temp_profile_file,
            )

            if _using_adiabat and prev_pressure is not None:
                # Bump blend toward full adiabat
                _adiabat_blend = min(1.0, _adiabat_blend + _ADIABAT_BLEND_STEP)
                verbose and logger.info(
                    f'Outer iter {outer_iter}: adiabat blend = {_adiabat_blend:.2f}'
                )

                # Recompute adiabat from previous iteration's converged structure
                adiabat_T = compute_adiabatic_temperature(
                    prev_radii,
                    prev_pressure,
                    prev_mass_enclosed,
                    surface_temperature,
                    cmb_mass,
                    core_mantle_mass,
                    layer_mixtures,
                    material_dictionaries,
                    interpolation_cache,
                    solidus_func,
                    liquidus_func,
                    mushy_zone_factors,
                    condensed_rho_min,
                    condensed_rho_scale,
                    binodal_T_scale,
                )

                # Build T(P) interpolator from the previous iteration's
                # pressure profile.  This ensures that during the Brent
                # bracket search, the adiabatic temperature tracks the
                # ODE's actual pressure rather than a fixed radial mapping.
                # This prevents unphysical (low P, high T) queries that
                # hit NaN gaps in the PALEOS tables.
                _sort = np.argsort(prev_pressure)
                _P_sorted = prev_pressure[_sort]
                _T_sorted = adiabat_T[_sort]
                # Remove P <= 0 entries (surface padding)
                _valid = _P_sorted > 0
                _P_sorted = _P_sorted[_valid]
                _T_sorted = _T_sorted[_valid]
                _logP_sorted = np.log10(_P_sorted)

                # Clamp adiabat T to a physically reasonable range.
                # np.interp flat-extrapolates at the edges, but the Brent
                # solver may query at extreme P values far beyond the
                # converged profile, producing huge T. Cap at 100,000 K
                # (PALEOS table maximum) and floor at 100 K.
                _T_MAX_CLAMP = 100000.0
                _T_MIN_CLAMP = 100.0

                if _adiabat_blend < 1.0:
                    _blend = _adiabat_blend

                    def temperature_function(
                        r, P, _b=_blend, _lp=_logP_sorted, _ts=_T_sorted, _ltf=_linear_tf
                    ):
                        T_lin = _ltf(r)
                        if P <= 0:
                            return T_lin
                        T_adi = float(np.interp(np.log10(P), _lp, _ts))
                        T_adi = max(_T_MIN_CLAMP, min(T_adi, _T_MAX_CLAMP))
                        return (1.0 - _b) * T_lin + _b * T_adi

                else:

                    def temperature_function(r, P, _lp=_logP_sorted, _ts=_T_sorted):
                        if P <= 0:
                            return surface_temperature
                        T_val = float(np.interp(np.log10(P), _lp, _ts))
                        return max(_T_MIN_CLAMP, min(T_val, _T_MAX_CLAMP))

                # Pre-compute temperatures array for the density update loop
                # (uses the converged pressure from the previous iteration)
                temperatures = np.array(
                    [
                        temperature_function(radii[i], prev_pressure[i])
                        for i in range(num_layers)
                    ]
                )
            else:

                def temperature_function(r, P, _tf=_linear_tf):
                    return _tf(r)

                temperatures = _linear_tf(radii)
        else:
            temperatures = np.ones(num_layers) * 300

        cmb_mass = core_mass_fraction * planet_mass
        core_mantle_mass = (core_mass_fraction + mantle_mass_fraction) * planet_mass

        pressure[0] = earth_center_pressure

        for inner_iter in range(max_iterations_inner):
            old_density = density.copy()

            # Scaling-law estimate for central pressure
            pressure_guess = (
                earth_center_pressure
                * (planet_mass / earth_mass) ** 2
                * (radius_guess / earth_radius) ** (-4)
            )
            # Cap the central pressure guess for WolfBower2018 (1 TPa table)
            # but not for RTPress100TPa (100 TPa melt table)
            uses_WB2018 = 'WolfBower2018:MgSiO3' in all_comps
            if uses_WB2018:
                pressure_guess = min(pressure_guess, max_center_pressure_guess)

            # Mutable state to capture the last ODE solution from inside
            # the residual function (brentq doesn't return intermediate
            # results, only the root)
            _state = {'mass_enclosed': None, 'gravity': None, 'pressure': None, 'n_evals': 0}

            def _pressure_residual(p_center):
                """Surface pressure residual f(P_c) = P_surface(P_c) - P_target.

                When the ODE integration terminates early (pressure hit zero
                before reaching the planet surface), P_center is too low and
                the residual is negative.
                """
                y0 = [0, 0, p_center]
                m, g, p = solve_structure(
                    layer_mixtures,
                    cmb_mass,
                    core_mantle_mass,
                    radii,
                    adaptive_radial_fraction,
                    relative_tolerance,
                    absolute_tolerance,
                    maximum_step,
                    material_dictionaries,
                    interpolation_cache,
                    y0,
                    solidus_func,
                    liquidus_func,
                    temperature_function if uses_Tdep else None,
                    mushy_zone_factors,
                    condensed_rho_min,
                    condensed_rho_scale,
                    binodal_T_scale,
                )
                if iteration_profiles_enabled:
                    create_pressure_density_files(
                        outer_iter, inner_iter, _state['n_evals'], radii, p, density
                    )
                _state['mass_enclosed'] = m
                _state['gravity'] = g
                _state['pressure'] = p
                _state['n_evals'] += 1

                # Early termination: pressure reached zero before the
                # surface (padded with zeros) → P_center is too low
                if p[-1] <= 0:
                    return -target_surface_pressure

                return p[-1] - target_surface_pressure

            # Bracket: P_center must be positive and wide enough to
            # straddle the root (where surface P = target P)
            p_low = max(1e6, 0.1 * pressure_guess)
            p_high = 10.0 * pressure_guess
            if uses_WB2018:
                p_high = min(p_high, max_center_pressure_guess)

            try:
                # Pre-validate that the bracket straddles the root.
                # This gives a clearer error than brentq's generic ValueError,
                # and the except handler below gracefully falls back to the
                # last evaluated solution.
                f_low = _pressure_residual(p_low)
                f_high = _pressure_residual(p_high)
                if f_low * f_high > 0:
                    raise ValueError(
                        f'Brent bracket does not straddle the root: '
                        f'f({p_low:.2e})={f_low:.2e}, f({p_high:.2e})={f_high:.2e}.'
                    )
                p_solution, root_info = brentq(
                    _pressure_residual,
                    p_low,
                    p_high,
                    xtol=1e6,
                    rtol=1e-10,
                    maxiter=max_iterations_pressure,
                    full_output=True,
                )
                # Re-run solve_structure at the exact root to get clean profiles
                # (brentq may have evaluated _state at a slightly different P)
                y0_root = [0, 0, p_solution]
                mass_enclosed, gravity, pressure = solve_structure(
                    layer_mixtures,
                    cmb_mass,
                    core_mantle_mass,
                    radii,
                    adaptive_radial_fraction,
                    relative_tolerance,
                    absolute_tolerance,
                    maximum_step,
                    material_dictionaries,
                    interpolation_cache,
                    y0_root,
                    solidus_func,
                    liquidus_func,
                    temperature_function if uses_Tdep else None,
                    mushy_zone_factors,
                    condensed_rho_min,
                    condensed_rho_scale,
                    binodal_T_scale,
                )

                surface_residual = abs(pressure[-1] - target_surface_pressure)
                # Allow zero pressure at the surface: the terminal event
                # pads truncated points with P=0, so check >= 0
                if (
                    root_info.converged
                    and surface_residual < pressure_tolerance
                    and np.min(pressure) >= 0
                ):
                    converged_pressure = True
                    verbose and logger.info(
                        f'Surface pressure converged after '
                        f'{root_info.function_calls} evaluations (Brent method).'
                    )
                else:
                    converged_pressure = False
                    verbose and logger.warning(
                        f'Brent method: converged={root_info.converged}, '
                        f'residual={surface_residual:.2e} Pa, '
                        f'min_P={np.min(pressure):.2e} Pa.'
                    )
            except ValueError:
                # f(p_low) and f(p_high) have the same sign — bracket
                # invalid.  Use the last evaluated solution if available.
                verbose and logger.warning(
                    f'Could not bracket pressure root in [{p_low:.2e}, {p_high:.2e}] Pa.'
                )
                if _state['mass_enclosed'] is not None:
                    mass_enclosed = _state['mass_enclosed']
                    gravity = _state['gravity']
                    pressure = _state['pressure']
                else:
                    # No evaluations succeeded — keep profiles from previous
                    # outer iteration (already initialised above).
                    verbose and logger.warning(
                        'No valid ODE solutions obtained during bracket search. '
                        'Keeping previous profiles.'
                    )
                converged_pressure = False

            # Update density grid using vectorized EOS lookups.
            # Partition shells by layer (all shells in a layer share one EOS),
            # then batch-call calculate_mixed_density_batch per layer.
            n_valid = min(num_layers, len(mass_enclosed))
            new_density = np.full(num_layers, np.nan)

            # Zero-pressure shells get zero density
            p_valid = pressure[:n_valid] > 0

            # Compute temperatures for valid shells
            if uses_Tdep:
                T_arr = np.array(
                    [
                        temperature_function(radii[i], pressure[i])
                        for i in range(n_valid)
                        if p_valid[i]
                    ]
                )
            else:
                T_arr = np.full(int(np.sum(p_valid)), 300.0)

            # Partition shells by layer using mass boundaries
            valid_indices = np.where(p_valid[:n_valid])[0]
            m_valid = mass_enclosed[valid_indices]
            rtol = 1e-12
            in_core = m_valid < cmb_mass * (1.0 - rtol)
            in_ice = ('ice_layer' in layer_mixtures) & (
                m_valid >= core_mantle_mass * (1.0 - rtol)
            )
            in_mantle = ~in_core & ~in_ice

            for layer_name, mask in [
                ('core', in_core),
                ('mantle', in_mantle),
                ('ice_layer', in_ice),
            ]:
                if not np.any(mask) or layer_name not in layer_mixtures:
                    continue
                idx = valid_indices[mask]
                mix = layer_mixtures[layer_name]
                rho_batch = calculate_mixed_density_batch(
                    pressure[idx],
                    T_arr[np.where(mask)[0]]
                    if len(T_arr) == len(valid_indices)
                    else T_arr[mask],
                    mix,
                    material_dictionaries,
                    solidus_func,
                    liquidus_func,
                    interpolation_cache,
                    mushy_zone_factors,
                    condensed_rho_min,
                    condensed_rho_scale,
                    binodal_T_scale,
                )
                new_density[idx] = rho_batch

            # Fill NaN entries with last valid density (walking outward)
            last_valid = None
            for i in range(n_valid):
                if not p_valid[i]:
                    new_density[i] = 0.0
                elif np.isnan(new_density[i]):
                    new_density[i] = last_valid if last_valid is not None else old_density[i]
                else:
                    last_valid = new_density[i]

            # Picard blend
            density[:n_valid] = 0.5 * (new_density[:n_valid] + old_density[:n_valid])

            # Check density convergence
            relative_diff_inner = np.max(
                np.abs((density - old_density) / (old_density + 1e-20))
            )
            if relative_diff_inner < tolerance_inner:
                verbose and logger.info(
                    f'Inner loop converged after {inner_iter + 1} iterations.'
                )
                converged_density = True
                break

            if inner_iter == max_iterations_inner - 1:
                verbose and logger.warning(
                    f'Maximum inner iterations ({max_iterations_inner}) reached. '
                    'Density may not be fully converged.'
                )

        # Recompute the temperatures array from the converged pressure profile
        # so model_results['temperature'] reflects actual T(P), not the pre-Brent estimate.
        if uses_Tdep:
            temperatures = np.array(
                [temperature_function(radii[i], pressure[i]) for i in range(num_layers)]
            )

        # Save converged profiles for the next outer iteration's adiabat
        prev_radii = radii.copy()
        prev_pressure = np.asarray(pressure).copy()
        prev_mass_enclosed = np.asarray(mass_enclosed).copy()

        # Update radius guess with damped scaling to prevent oscillation.
        # The cube-root scaling is correct in direction but can overshoot
        # wildly when calculated_mass << planet_mass, catapulting radius
        # to unphysical values and trapping the solver in a cycle.
        calculated_mass = mass_enclosed[-1]
        if calculated_mass <= 0 or not np.isfinite(calculated_mass):
            radius_guess *= 0.8
            verbose and logger.warning(
                f'Outer iter {outer_iter}: calculated_mass={calculated_mass:.2e}, '
                f'shrinking radius_guess to {radius_guess:.0f} m.'
            )
        else:
            scale = (planet_mass / calculated_mass) ** (1.0 / 3.0)
            scale = max(0.5, min(scale, 2.0))
            radius_guess *= scale
        cmb_mass = core_mass_fraction * calculated_mass
        core_mantle_mass = (core_mass_fraction + mantle_mass_fraction) * calculated_mass

        relative_diff_outer_mass = np.abs((calculated_mass - planet_mass) / planet_mass)

        # MASS CONVERGENCE CHECK
        # When temperature_mode='adiabatic' and the blend has not yet reached
        # 1.0, mass convergence triggers the adiabat transition instead of
        # breaking. The blend ramps 0 -> 0.5 -> 1.0 over successive mass
        # convergences, preventing solver divergence.
        if relative_diff_outer_mass < tolerance_outer:
            if temperature_mode == 'adiabatic' and _adiabat_blend < 1.0:
                if not _using_adiabat:
                    _using_adiabat = True
                    logger.info(
                        f'Outer iter {outer_iter}: mass converged with linear T, '
                        f'activating adiabat blend.'
                    )
                # Continue iterating to let the blend ramp up
                continue
            logger.info(f'Outer loop (total mass) converged after {outer_iter + 1} iterations.')
            converged_mass = True
            break

        if outer_iter == max_iterations_outer - 1:
            verbose and logger.warning(
                f'Maximum outer iterations ({max_iterations_outer}) reached. '
                'Total mass may not be fully converged.'
            )

    if converged_mass and converged_density and converged_pressure:
        converged = True

    end_time = time.time()
    total_time = end_time - start_time

    # Clean up surface artifacts. The Picard iteration can produce
    # density jumps at the outermost shells where the pressure drops
    # rapidly and the EOS lookup is sensitive to small P changes.
    # Detect shells where the density gradient suddenly steepens
    # (more than 3x the running average) and replace them with a
    # linear extrapolation from the smooth interior.
    pressure = np.asarray(pressure)
    density = np.asarray(density)

    # Zero out padded (P=0) shells
    density[pressure <= 0] = 0.0

    # Find the last shell with positive density
    i_surf = len(density) - 1
    while i_surf > 0 and density[i_surf] <= 0:
        i_surf -= 1

    if i_surf > 10:
        # Compute density gradient in the outer mantle.
        # Use the last 20 shells before the surface edge, but do not
        # cross layer boundaries (CMB, ice-layer transition) where
        # real density jumps exist.
        # Find outermost layer boundary by detecting large density
        # ratios (> 1.5x) walking inward from surface.
        i_layer_base = max(0, i_surf - 40)
        for _k in range(i_surf - 1, i_layer_base, -1):
            if density[_k] > 0 and density[_k + 1] > 0:
                ratio = density[_k] / density[_k + 1]
                if ratio > 1.5 or ratio < 1.0 / 1.5:
                    i_layer_base = _k + 1
                    break
        n_check = min(20, i_surf - i_layer_base)
        if n_check < 6:
            n_check = min(20, i_surf - 5)
        i_start = i_surf - n_check
        grads = np.diff(density[i_start : i_surf + 1])
        # Find where the gradient suddenly changes
        # (the smooth interior has nearly constant negative gradient).
        if len(grads) > 5:
            median_grad = np.median(grads[: n_check // 2])
            if median_grad < 0:  # density decreasing outward (normal)
                for j in range(len(grads)):
                    i_global = i_start + j + 1
                    # Catch both sudden drops (3x steeper) and sudden
                    # increases (sign reversal). Both indicate the Picard
                    # iteration produced unreliable density at the surface.
                    is_sudden_drop = grads[j] < 3 * median_grad
                    is_sudden_rise = grads[j] > 0 and abs(grads[j]) > abs(median_grad)
                    if is_sudden_drop or is_sudden_rise:
                        # Extrapolate linearly from the smooth region
                        for k in range(i_global, i_surf + 1):
                            extrap = density[i_global - 1] + median_grad * (k - i_global + 1)
                            density[k] = max(extrap, 0.0)
                        break

    model_results = {
        'layer_eos_config': layer_eos_config,
        'radii': radii,
        'density': density,
        'gravity': gravity,
        'pressure': pressure,
        'temperature': temperatures,
        'mass_enclosed': mass_enclosed,
        'cmb_mass': cmb_mass,
        'core_mantle_mass': core_mantle_mass,
        'total_time': total_time,
        'converged': converged,
        'converged_pressure': converged_pressure,
        'converged_density': converged_density,
        'converged_mass': converged_mass,
    }
    return model_results

def post_processing(config_params, id_mass=None, output_file=None):
    """Post-process model results by saving output data and plotting.

    Parameters
    ----------
    config_params : dict
        Configuration parameters for the model.
    id_mass : str or None
        Identifier for the planet mass, used in output file naming.
    output_file : str or None
        Path to the output file for calculated mass and radius.
    """
    data_output_enabled = config_params['data_output_enabled']
    plotting_enabled = config_params['plotting_enabled']

    layer_eos_config = config_params['layer_eos_config']
    solidus_id = config_params.get('rock_solidus', 'Stixrude14-solidus')
    liquidus_id = config_params.get('rock_liquidus', 'Stixrude14-liquidus')

    model_results = main(
        config_params,
        material_dictionaries=load_material_dictionaries(),
        melting_curves_functions=load_solidus_liquidus_functions(
            layer_eos_config, solidus_id, liquidus_id
        ),
        input_dir=os.path.join(ZALMOXIS_ROOT, 'input'),
    )

    # Extract results
    radii = model_results['radii']
    density = model_results['density']
    gravity = model_results['gravity']
    pressure = model_results['pressure']
    temperature = model_results['temperature']
    mass_enclosed = model_results['mass_enclosed']
    cmb_mass = model_results['cmb_mass']
    core_mantle_mass = model_results['core_mantle_mass']
    total_time = model_results['total_time']
    converged = model_results['converged']
    converged_pressure = model_results['converged_pressure']
    converged_density = model_results['converged_density']
    converged_mass = model_results['converged_mass']

    cmb_index = np.argmax(mass_enclosed >= cmb_mass)

    average_density = mass_enclosed[-1] / (4 / 3 * math.pi * radii[-1] ** 3)

    # Check if mantle uses a Tdep EOS that needs external melting curves
    # for phase detection. Unified PALEOS tables derive phases from the
    # table itself and do not need (or support) get_Tdep_material().
    mantle_str = layer_eos_config.get('mantle', '')
    if mantle_str:
        _mantle_mix = parse_layer_components(mantle_str)
        uses_phase_detection = bool(set(_mantle_mix.components) & _NEEDS_MELTING_CURVES)
    else:
        uses_phase_detection = False

    if uses_phase_detection:
        mantle_pressures = pressure[cmb_index:]
        mantle_temperatures = temperature[cmb_index:]
        mantle_radii = radii[cmb_index:]

        solidus_func, liquidus_func = load_solidus_liquidus_functions(
            layer_eos_config, solidus_id, liquidus_id
        )

        mantle_phases = get_Tdep_material(
            mantle_pressures, mantle_temperatures, solidus_func, liquidus_func
        )

    logger.info('Exoplanet Internal Structure Model Results:')
    logger.info('----------------------------------------------------------------------')
    logger.info(
        f'Calculated Planet Mass: {mass_enclosed[-1]:.2e} kg or '
        f'{mass_enclosed[-1] / earth_mass:.2f} Earth masses'
    )
    logger.info(
        f'Calculated Planet Radius: {radii[-1]:.2e} m or '
        f'{radii[-1] / earth_radius:.2f} Earth radii'
    )
    logger.info(f'Core Radius: {radii[cmb_index]:.2e} m')
    logger.info(f'Mantle Density (at CMB): {density[cmb_index]:.2f} kg/m^3')
    logger.info(f'Core Density (at CMB): {density[cmb_index - 1]:.2f} kg/m^3')
    logger.info(f'Pressure at Core-Mantle Boundary (CMB): {pressure[cmb_index]:.2e} Pa')
    logger.info(f'Pressure at Center: {pressure[0]:.2e} Pa')
    logger.info(f'Average Density: {average_density:.2f} kg/m^3')
    logger.info(f'CMB Mass Fraction: {mass_enclosed[cmb_index] / mass_enclosed[-1]:.3f}')
    logger.info(
        f'Core+Mantle Mass Fraction: '
        f'{(core_mantle_mass - mass_enclosed[cmb_index]) / mass_enclosed[-1]:.3f}'
    )
    logger.info(f'Calculated Core Radius Fraction: {radii[cmb_index] / radii[-1]:.2f}')
    logger.info(
        f'Calculated Core+Mantle Radius Fraction: '
        f'{(radii[np.argmax(mass_enclosed >= core_mantle_mass)] / radii[-1]):.2f}'
    )
    logger.info(f'Total Computation Time: {total_time:.2f} seconds')
    logger.info(
        f'Overall Convergence Status: {converged} with Pressure: {converged_pressure}, '
        f'Density: {converged_density}, Mass: {converged_mass}'
    )

    if data_output_enabled:
        output_data = np.column_stack(
            (radii, density, gravity, pressure, temperature, mass_enclosed)
        )
        header = (
            'Radius (m)\tDensity (kg/m^3)\tGravity (m/s^2)\t'
            'Pressure (Pa)\tTemperature (K)\tMass Enclosed (kg)'
        )
        if id_mass is None:
            np.savetxt(
                os.path.join(ZALMOXIS_ROOT, 'output_files', 'planet_profile.txt'),
                output_data,
                header=header,
            )
        else:
            np.savetxt(
                os.path.join(ZALMOXIS_ROOT, 'output_files', f'planet_profile{id_mass}.txt'),
                output_data,
                header=header,
            )
        if output_file is None:
            output_file = os.path.join(
                ZALMOXIS_ROOT, 'output_files', 'calculated_planet_mass_radius.txt'
            )
        if not os.path.exists(output_file):
            header = 'Calculated Mass (kg)\tCalculated Radius (m)'
            with open(output_file, 'w') as file:
                file.write(header + '\n')
        with open(output_file, 'a') as file:
            file.write(f'{mass_enclosed[-1]}\t{radii[-1]}\n')

    if plotting_enabled:
        plot_planet_profile_single(
            radii,
            density,
            gravity,
            pressure,
            temperature,
            radii[np.argmax(mass_enclosed >= cmb_mass)],
            cmb_mass,
            mass_enclosed[-1] / (4 / 3 * math.pi * radii[-1] ** 3),
            mass_enclosed,
            id_mass,
            layer_eos_config=layer_eos_config,
        )

        if uses_phase_detection:
            plot_PT_with_phases(
                mantle_pressures,
                mantle_temperatures,
                mantle_radii,
                mantle_phases,
                radii[cmb_index],
            )