"""
Compute single physical fields.

These functions are high-level wrappers around C-functions that compute 3D fields, for
example initial conditions, perturbed fields and ionization fields.
"""

import logging
import warnings

import numpy as np
from astropy import units as un
from astropy.cosmology import z_at_value

from ..wrapper.arrays import Array
from ..wrapper.inputs import InputParameters
from ..wrapper.outputs import (
    BrightnessTemp,
    HaloBox,
    HaloCatalog,
    InitialConditions,
    IonizedBox,
    PerturbedField,
    PerturbedHaloCatalog,
    TsBox,
    XraySourceBox,
)
from ._param_config import (
    check_output_consistency,
    single_field_func,
)

logger = logging.getLogger(__name__)


@single_field_func
def compute_initial_conditions(
    *, inputs: InputParameters, initial_density: np.ndarray | float | None = None
) -> InitialConditions:
    r"""
    Compute initial conditions.

    Parameters
    ----------
    inputs
        The InputParameters instance defining the run.
    initial_density: np.ndarray or float, optional
        A realization of the density field on the high resolution grid.
        This input can also be used to determine the global density field,
        in case we have a single cell in the box.
    regenerate : bool, optional
        Whether to force regeneration of data, even if matching cached data is found.
    cache
        An OutputCache object defining how to read cached boxes.
    write
        A boolean specifying whether we need to cache the box.

    Returns
    -------
    :class:`~InitialConditions`
    """
    # Initialize memory for the boxes that will be returned.
    ics = InitialConditions.new(inputs=inputs)

    if inputs.simulation_options.HII_DIM == 1 and inputs.simulation_options.DIM == 1:
        # If we have only one cell (could happen if we do a global evolution), we don't really need to compute the ics
        shape = (1, 1, 1)
        required_arrays = PerturbedField.new(
            redshift=0, inputs=inputs
        ).get_required_input_arrays(ics)

        # Set the arrays to zero, or according to initial_density if the arrays are density fields
        for array in required_arrays:
            if (initial_density is not None) and (
                array in ["hires_density", "lowres_density"]
            ):
                value = initial_density
            else:
                value = 0.0
            setattr(
                ics,
                array,
                Array(shape=shape, dtype=np.float32)
                .initialize()
                .with_value(val=value * np.ones(shape)),
            )
        return ics
    else:
        if initial_density is not None:
            if np.abs(initial_density.mean() > 1e-3):
                warnings.warn(
                    f"Your initial_density has mean {initial_density.mean()}. "
                    + "Make sure you know what you are doing.",
                    stacklevel=2,
                )
            shape = ics.hires_density.shape
            if initial_density.shape != shape:
                raise ValueError(
                    "The shape of your high resolution initial_density is not consistent with inputs!"
                    + f" According to inputs, initial_density must be of shape {shape}, got {initial_density.shape}."
                )

            ics.hires_density = (
                Array(shape=shape, dtype=np.float32)
                .initialize()
                ._with_value_not_computed(val=initial_density)
            )
        return ics.compute()


@single_field_func
def perturb_field(
    *,
    redshift: float,
    inputs: InputParameters | None = None,
    initial_conditions: InitialConditions,
) -> PerturbedField:
    r"""
    Compute a perturbed field at a given redshift.

    Parameters
    ----------
    redshift : float
        The redshift at which to compute the perturbed field.
    initial_conditions : :class:`~InitialConditions` instance
        The initial conditions.

    Returns
    -------
    :class:`~PerturbedField`

    Other Parameters
    ----------------
    regenerate, write, cache:
        See docs of :func:`initial_conditions` for more information.

    Examples
    --------
    >>> initial_conditions = compute_initial_conditions()
    >>> field7 = perturb_field(7.0, initial_conditions)
    >>> field8 = perturb_field(8.0, initial_conditions)

    The user and cosmo parameter structures are by default inferred from the
    ``initial_conditions``.
    """
    # Initialize perturbed boxes.
    fields = PerturbedField.new(redshift=redshift, inputs=inputs)

    # Run the C Code
    return fields.compute(ics=initial_conditions)


@single_field_func
def determine_halo_catalog(
    *,
    redshift: float,
    inputs: InputParameters | None = None,
    initial_conditions: InitialConditions,
    descendant_halos: HaloCatalog | None = None,
) -> HaloCatalog:
    r"""
    Find a halo list, given a redshift.

    Parameters
    ----------
    redshift : float
        The redshift at which to determine the halo list.
    initial_conditions : :class:`~InitialConditions` instance
        The initial conditions fields (density, velocity).
    descendant_halos : :class:`~HaloCatalog` instance, optional
        The halos that form the descendants (i.e. lower redshift) of those computed by
        this function. If this is not provided, we generate the initial stochastic halos
        directly in this function (and progenitors can then be determined by these).

    Returns
    -------
    :class:`~HaloCatalog`

    Other Parameters
    ----------------
    regenerate, write, cache:
        See docs of :func:`initial_conditions` for more information.
    """
    if inputs.matter_options.HMF != "ST":
        warnings.warn(
            "DexM Halofinder sses a fit to the Sheth-Tormen mass function."
            "With HMF!=1 the Halos from DexM will not be from the same mass function",
            stacklevel=2,
        )

    if descendant_halos is None:
        descendant_halos = HaloCatalog.dummy()

    # Initialize halo list boxes.
    fields = HaloCatalog.new(
        redshift=redshift,
        desc_redshift=descendant_halos.redshift,
        inputs=inputs,
    )

    # Run the C Code
    return fields.compute(
        ics=initial_conditions,
        descendant_halos=descendant_halos,
    )


@single_field_func
def perturb_halo_catalog(
    *,
    initial_conditions: InitialConditions,
    inputs: InputParameters | None = None,
    previous_spin_temp: TsBox | None = None,
    previous_ionize_box: IonizedBox | None = None,
    halo_catalog: HaloCatalog,
) -> PerturbedHaloCatalog:
    r"""
    Given a halo list, perturb the halos for a given redshift.

    Parameters
    ----------
    initial_conditions : :class:`~InitialConditions`
        The initial conditions of the run. The user and cosmo params
        as well as the random seed will be set from this object.
    halo_catalog: :class: `~HaloCatalog`
        The halo catalogue in Lagrangian space to be perturbed.

    Returns
    -------
    :class:`~PerturbedHaloCatalog`

    Other Parameters
    ----------------
    regenerate, write, direc:
        See docs of :func:`initial_conditions` for more information.

    Examples
    --------
    Fill this in once finalised

    """
    hbuffer_size = (
        halo_catalog.n_halos if halo_catalog.n_halos else halo_catalog.buffer_size
    )
    redshift = halo_catalog.redshift

    # Initialize halo list boxes.
    fields = PerturbedHaloCatalog.new(
        redshift=redshift,
        buffer_size=hbuffer_size,
        inputs=inputs,
    )
    if previous_spin_temp is None:
        if (
            redshift >= inputs.simulation_options.Z_HEAT_MAX
            or not inputs.astro_options.USE_MINI_HALOS
        ):
            # Dummy spin temp is OK since we're above Z_HEAT_MAX
            previous_spin_temp = TsBox.dummy()
        else:
            raise ValueError("Below Z_HEAT_MAX you must specify the previous_spin_temp")

    if previous_ionize_box is None:
        if (
            redshift >= inputs.simulation_options.Z_HEAT_MAX
            or not inputs.astro_options.USE_MINI_HALOS
        ):
            # Dummy ionize box is OK since we're above Z_HEAT_MAX
            previous_ionize_box = IonizedBox.dummy()
        else:
            raise ValueError(
                "Below Z_HEAT_MAX you must specify the previous_ionize_box"
            )

    # Run the C Code
    return fields.compute(
        ics=initial_conditions,
        halo_catalog=halo_catalog,
        previous_spin_temp=previous_spin_temp,
        previous_ionize_box=previous_ionize_box,
    )


@single_field_func
def compute_halo_grid(
    *,
    redshift: float,
    initial_conditions: InitialConditions,
    inputs: InputParameters | None = None,
    halo_catalog: HaloCatalog | None = None,
    previous_spin_temp: TsBox | None = None,
    previous_ionize_box: IonizedBox | None = None,
) -> HaloBox:
    r"""
    Compute grids of halo properties from a catalogue.

    At the moment this simply produces halo masses, stellar masses and SFR on a grid of
    HII_DIM. In the future this will compute properties such as emissivities which will
    be passed directly into ionize_box etc. instead of the catalogue.

    Parameters
    ----------
    initial_conditions : :class:`~InitialConditions`
        The initial conditions of the run.
    inputs : :class:`~InputParameters`, optional
        The input parameters specifying the run.
    halo_catalog: :class:`~HaloCatalog`, optional
        This contains all the dark matter haloes obtained if using a discrete halo model.
        This is a list of halo masses and coordinates for the dark matter halos.
    previous_spin_temp : :class:`TsBox`, optional
        The previous spin temperature box. Used for feedback when USE_MINI_HALOS==True
    previous_ionize_box: :class:`IonizedBox` or None
        An at the last timestep. Used for feedback when USE_MINI_HALOS==True

    Returns
    -------
    :class:`~HaloBox` :
        An object containing the halo box data.

    Other Parameters
    ----------------
    regenerate, write, cache:
        See docs of :func:`initial_conditions` for more information.
    """
    box = HaloBox.new(redshift=redshift, inputs=inputs)

    if halo_catalog is None:
        if inputs.matter_options.has_discrete_halos:
            raise ValueError(
                f"You must provide halo_catalog for SOURCE_MODEL = {inputs.matter_options.SOURCE_MODEL}"
            )
        else:
            halo_catalog = HaloCatalog.dummy()

    # NOTE: due to the order, we use the previous spin temp here, like spin_temperature,
    #       but UNLIKE ionize_box, which uses the current box
    # TODO: think about the inconsistency here
    # NOTE: if USE_MINI_HALOS is TRUE, so is USE_TS_FLUCT and INHOMO_RECO
    if previous_spin_temp is None:
        if (
            redshift >= inputs.simulation_options.Z_HEAT_MAX
            or not inputs.astro_options.USE_MINI_HALOS
        ):
            # Dummy spin temp is OK since we're above Z_HEAT_MAX
            previous_spin_temp = TsBox.dummy()
        else:
            raise ValueError("Below Z_HEAT_MAX you must specify the previous_spin_temp")

    if previous_ionize_box is None:
        if (
            redshift >= inputs.simulation_options.Z_HEAT_MAX
            or not inputs.astro_options.USE_MINI_HALOS
        ):
            # Dummy ionize box is OK since we're above Z_HEAT_MAX
            previous_ionize_box = IonizedBox.dummy()
        else:
            raise ValueError(
                "Below Z_HEAT_MAX you must specify the previous_ionize_box"
            )

    return box.compute(
        initial_conditions=initial_conditions,
        halo_catalog=halo_catalog,
        previous_ionize_box=previous_ionize_box,
        previous_spin_temp=previous_spin_temp,
    )


# TODO: make this more general and probably combine with the lightcone interp function
def interp_halo_boxes(
    halo_boxes: list[HaloBox],
    fields: list[str],
    redshift: float,
) -> HaloBox:
    """
    Interpolate HaloBox history to the desired redshift.

    Photon conservation & Xray sources require halo boxes at redshifts
    that are not equal to the current redshift, and may be between redshift steps.
    So we need a function to interpolate between two halo boxes.
    We assume here that z_arr is strictly INCERASING

    Parameters
    ----------
    halo_boxes : list of HaloBox instances
        The halobox history to be interpolated
    fields: List of Strings
        The properties of the haloboxes to be interpolated
    redshift : float
        The desired redshift of interpolation

    Returns
    -------
    :class:`~HaloBox` :
        An object containing the halo box data
    """
    inputs = halo_boxes[0].inputs
    z_halos = [box.redshift for box in halo_boxes]
    if not np.all(np.diff(z_halos) > 0):
        raise ValueError("halo_boxes must be in ascending order of redshift")

    if redshift > z_halos[-1] or redshift < z_halos[0]:
        raise ValueError(f"Invalid z_target {redshift} for redshift array {z_halos}")

    # If we do global evolution, no need to do that
    if inputs.simulation_options.HII_DIM > 1:
        arr_fields = [f for f in fields if f in halo_boxes[0].arrays]
        computed = [box.ensure_arrays_computed(*arr_fields) for box in halo_boxes]
        if not all(computed):
            raise ValueError("Some of the HaloBox fields required are not computed")

    idx_prog = np.searchsorted(z_halos, redshift, side="left")

    if idx_prog == 0 or idx_prog == len(z_halos):
        logger.debug(f"redshift {redshift} beyond limits, {z_halos[0], z_halos[-1]}")
        raise ValueError

    z_prog = z_halos[idx_prog]
    idx_desc = idx_prog - 1

    z_desc = z_halos[idx_desc]
    interp_param = (redshift - z_desc) / (z_prog - z_desc)

    # If we do global evolution, no need to do that
    if inputs.simulation_options.HII_DIM > 1:
        # I set the box redshift to be the stored one so it is read properly into the ionize box
        # for the xray source it doesn't matter, also since it is not _compute()'d, it won't be cached
        check_output_consistency(
            dict(
                zip(
                    [f"box-{i}" for i in range(len(halo_boxes))],
                    halo_boxes,
                    strict=True,
                )
            )
        )
    hbox_out = HaloBox.new(redshift=redshift, inputs=inputs)

    # initialise the memory
    hbox_out._init_arrays()

    # interpolate halo boxes in gridded SFR
    hbox_prog = halo_boxes[idx_prog]
    hbox_desc = halo_boxes[idx_desc]

    for field in fields:
        field_desc = hbox_desc.get(field)
        field_prog = hbox_prog.get(field)
        interp_field = np.zeros_like(field_desc)
        interp_field[...] = (1 - interp_param) * field_desc + interp_param * field_prog
        hbox_out.set(field, interp_field)

    return hbox_out


# NOTE: the current implementation of this box is very hacky, since I have trouble figuring out a way to _compute()
#   over multiple redshifts in a nice way using this wrapper.
# TODO: if we move some code to jax or similar I think this would be one of the first candidates (just filling out some filtered grids)
@single_field_func
def compute_xray_source_field(
    *,
    initial_conditions: InitialConditions,
    hboxes: list[HaloBox],
    redshift: float,
) -> XraySourceBox:
    r"""
    Compute filtered grid of SFR for use in spin temperature calculation.

    This will filter over the halo history in annuli, computing the contribution to the
    SFR density

    If no halo field is passed one is calculated at the desired redshift as if it is the
    first box.

    Parameters
    ----------
    initial_conditions : :class:`~InitialConditions`
        The initial conditions of the run. The user and cosmo params
    hboxes: Sequence of :class:`~HaloBox` instances
        This contains the list of Halobox instances which are used to create this source field

    Returns
    -------
    :class:`~XraySourceBox` :
        An object containing x ray heating, ionisation, and lyman alpha rates.

    Other Parameters
    ----------------
    regenerate, write, cache:
        See docs of :func:`initial_conditions` for more information.
    """
    z_halos = [hb.redshift for hb in hboxes]
    inputs = hboxes[0].inputs

    # Initialize halo list boxes.
    box = XraySourceBox.new(redshift=redshift, inputs=inputs)

    # set minimum R at cell size
    l_factor = (4 * np.pi / 3.0) ** (-1 / 3)
    if inputs.simulation_options.HII_DIM == 1:
        # If HII_DIM=1 (happens when we run_global_evolution), we take a typical cell size of 1.5Mpc,
        # just to for setting the z'' array (note that filtering won't be done on a box with a single cell)
        R_min = 1.5 * l_factor
    else:
        R_min = (
            inputs.simulation_options.BOX_LEN
            / inputs.simulation_options.HII_DIM
            * l_factor
        )
    z_max = min(max(z_halos), inputs.simulation_options.Z_HEAT_MAX)

    # now we need to find the closest halo box to the redshift of the shell
    cosmo_ap = inputs.cosmo_params.cosmo
    cmd_zp = cosmo_ap.comoving_distance(redshift)
    R_steps = np.arange(0, inputs.astro_params.N_STEP_TS)
    R_factor = (inputs.astro_params.R_MAX_TS / R_min) ** (
        R_steps / inputs.astro_params.N_STEP_TS
    )
    R_range = un.Mpc * R_min * R_factor
    cmd_edges = cmd_zp + R_range  # comoving distance edges
    # Get the edges of the shells
    zmin = z_at_value(cosmo_ap.comoving_distance, cmd_edges.min()).value
    zmax = z_at_value(cosmo_ap.comoving_distance, cmd_edges.max()).value
    zgrid = np.logspace(np.log10(zmin), np.log10(zmax), 100)
    dgrid = cosmo_ap.comoving_distance(zgrid)
    zpp_edges = np.interp(cmd_edges.value, dgrid.value, zgrid)
    # the `average` redshift of the shell is the average of the
    # inner and outer redshifts (following the C code)
    zpp_avg = zpp_edges - np.diff(np.insert(zpp_edges, 0, redshift)) / 2

    interp_fields = ["halo_sfr", "halo_xray"]
    if inputs.astro_options.USE_MINI_HALOS:
        interp_fields += ["halo_sfr_mini", "log10_Mcrit_MCG_ave"]

    # call the box the initialize the memory, since I give some values before computing
    box._init_arrays()
    for i in range(inputs.astro_params.N_STEP_TS):
        R_inner = R_range[i - 1].to("Mpc").value if i > 0 else 0
        R_outer = R_range[i].to("Mpc").value

        if zpp_avg[i] >= z_max:
            box.filtered_sfr.value[i] = 0
            box.filtered_xray.value[i] = 0
            if inputs.astro_options.USE_MINI_HALOS:
                box.filtered_sfr_mini.value[i] = 0
                box.mean_log10_Mcrit_LW.value[i] = inputs.astro_params.M_TURN  # minimum
            logger.debug(f"ignoring Radius {i} which is above Z_HEAT_MAX")
            continue

        hbox_interp = interp_halo_boxes(
            halo_boxes=hboxes[::-1],
            fields=interp_fields,
            redshift=zpp_avg[i],
        )

        # if we have no halos we ignore the whole shell
        sfr_allzero = np.all(hbox_interp.get("halo_sfr") == 0)
        if inputs.astro_options.USE_MINI_HALOS:
            sfr_allzero = sfr_allzero & np.all(hbox_interp.get("halo_sfr_mini") == 0)
        if sfr_allzero:
            box.filtered_sfr.value[i] = 0
            box.filtered_xray.value[i] = 0
            if inputs.astro_options.USE_MINI_HALOS:
                box.filtered_sfr_mini.value[i] = 0
                box.mean_log10_Mcrit_LW.value[i] = hbox_interp.log10_Mcrit_MCG_ave
            logger.debug(f"ignoring Radius {i} due to no stars")
            continue

        box = box.compute(
            halobox=hbox_interp,
            R_inner=R_inner,
            R_outer=R_outer,
            R_ct=i,
            allow_already_computed=True,
        )

    # Sometimes we don't compute at all
    # (if the first zpp > z_max or there are no halos at max R)
    # in which case the array is not marked as computed
    if not box.is_computed:
        for name, array in box.arrays.items():
            setattr(box, name, array.computed())

    return box


@single_field_func
def compute_spin_temperature(
    *,
    initial_conditions: InitialConditions,
    perturbed_field: PerturbedField,
    inputs: InputParameters | None = None,
    xray_source_box: XraySourceBox | None = None,
    previous_spin_temp: TsBox | None = None,
    cleanup: bool = False,
) -> TsBox:
    r"""
    Compute spin temperature boxes at a given redshift.

    See the notes below for how the spin temperature field is evolved through redshift.

    Parameters
    ----------
    initial_conditions : :class:`~InitialConditions`
        The initial conditions
    inputs : :class:`~InputParameters`
        The input parameters specifying the run. Since this will be the first box
        to use the astro params/flags when SOURCE_MODEL='E-INTEGRAL' and USE_TS_FLUCT=True.
    perturbed_field : :class:`~PerturbedField`, optional
        If given, this field will be used, otherwise it will be generated. To be generated,
        either `initial_conditions` and `redshift` must be given, or `simulation_options`, `cosmo_params` and
        `redshift`. By default, this will be generated at the same redshift as the spin temperature
        box. The redshift of perturb field is allowed to be different than `redshift`. If so, it
        will be interpolated to the correct redshift, which can provide a speedup compared to
        actually computing it at the desired redshift.
    xray_source_box : :class:`XraySourceBox`, optional
        When using a lagrangian source model, this box specifies the filtered sfr and xray emissivity at all
        redshifts/filter radii required by the spin temperature algorithm.
    previous_spin_temp : :class:`TsBox` or None
        The previous spin temperature box. Needed when we are beyond the first snapshot

    Returns
    -------
    :class:`~TsBox`
        An object containing the spin temperature box data.

    Other Parameters
    ----------------
    regenerate, write, cache:
        See docs of :func:`initial_conditions` for more information.
    """
    redshift = perturbed_field.redshift

    if redshift >= inputs.simulation_options.Z_HEAT_MAX:
        previous_spin_temp = TsBox.new(inputs=inputs, redshift=0.0, dummy=True)

    if xray_source_box is None:
        if inputs.matter_options.lagrangian_source_grid:
            raise ValueError(
                f"xray_source_box is required for SOURCE_MODEL= {inputs.matter_options.SOURCE_MODEL}"
            )
        else:
            xray_source_box = XraySourceBox.dummy()

    # Set up the box without computing anything.
    box = TsBox.new(
        redshift=redshift,
        inputs=inputs,
    )

    # Run the C Code
    return box.compute(
        cleanup=cleanup,
        perturbed_field=perturbed_field,
        xray_source_box=xray_source_box,
        prev_spin_temp=previous_spin_temp,
        ics=initial_conditions,
    )


@single_field_func
def compute_ionization_field(
    *,
    perturbed_field: PerturbedField,
    initial_conditions: InitialConditions,
    inputs: InputParameters | None = None,
    previous_perturbed_field: PerturbedField | None = None,
    previous_ionized_box: IonizedBox | None = None,
    spin_temp: TsBox | None = None,
    halobox: HaloBox | None = None,
) -> IonizedBox:
    r"""
    Compute an ionized box at a given redshift.

    This function has various options for how the evolution of the ionization is
    computed (if at all). See the Notes below for details.

    Parameters
    ----------
    initial_conditions : :class:`~InitialConditions` instance
        The initial conditions.
    inputs : :class:`~InputParameters`
        The input parameters specifying the run. Since this may be the first box
        to use the astro params/flags, it is needed when we have not computed a TsBox or HaloBox.
    perturbed_field : :class:`~PerturbedField`
        The perturbed density field.
    previous_perturbed_field : :class:`~PerturbedField`, optional
        An perturbed field at higher redshift. This is only used if USE_MINI_HALOS is included.
    previous_ionize_box: :class:`IonizedBox` or None
        An ionized box at higher redshift. This is only used if `INHOMO_RECO` and/or `USE_TS_FLUCT`
        are true. If either of these are true, and this is not given, then it will be assumed that
        this is the "first box", i.e. that it can be populated accurately without knowing source
        statistics.
    spin_temp: :class:`TsBox` or None, optional
        A spin-temperature box, only required if `USE_TS_FLUCT` is True. If None, will try to read
        in a spin temp box at the current redshift, and failing that will try to automatically
        create one, using the previous ionized box redshift as the previous spin temperature
        redshift.
    halobox: :class:`~HaloBox` or None, optional
        If passed, this contains all the dark matter haloes obtained if using the
        lagrangian source models. These are grids containing summed halo properties
        such as ionizing emissivity.

    Returns
    -------
    :class:`~IonizedBox` :
        An object containing the ionized box data.

    Notes
    -----
    Typically, the ionization field at any redshift is dependent on the evolution of xHI up until
    that redshift, which necessitates providing a previous ionization field to define the current
    one. If neither the spin temperature field, nor inhomogeneous recombinations (specified in
    flag options) are used, no evolution needs to be done. If the redshift is beyond
    Z_HEAT_MAX, previous fields are not required either.
    """
    redshift = perturbed_field.redshift

    if redshift >= inputs.simulation_options.Z_HEAT_MAX:
        # Previous boxes must be "initial"
        previous_ionized_box = IonizedBox.initial(inputs=inputs)
        previous_perturbed_field = PerturbedField.initial(inputs=inputs)

    if inputs.evolution_required:
        if previous_ionized_box is None:
            raise ValueError(
                "You need to provide a previous ionized box when redshift < Z_HEAT_MAX."
            )
        if previous_perturbed_field is None:
            raise ValueError(
                "You need to provide a previous perturbed field when redshift < Z_HEAT_MAX."
            )
    else:
        if previous_ionized_box is None:
            previous_ionized_box = IonizedBox.initial(inputs=inputs)
        if previous_perturbed_field is None:
            previous_perturbed_field = PerturbedField.initial(inputs=inputs)

    if inputs.simulation_options.HII_DIM > 1:
        box = IonizedBox.new(inputs=inputs, redshift=redshift)

        if not inputs.matter_options.lagrangian_source_grid:
            # Construct an empty halo field to pass in to the function.
            halobox = HaloBox.dummy()
        elif halobox is None:
            raise ValueError(
                f"A HaloBox must be provided for SOURCE_MODEL={inputs.matter_options.SOURCE_MODEL}"
            )

        # Set empty spin temp box if necessary.
        if not inputs.astro_options.USE_TS_FLUCT:
            spin_temp = TsBox.dummy()
        elif spin_temp is None:
            raise ValueError("No spin temperature box given but USE_TS_FLUCT=True")

        # Run the C Code
        return box.compute(
            perturbed_field=perturbed_field,
            prev_perturbed_field=previous_perturbed_field,
            prev_ionize_box=previous_ionized_box,
            spin_temp=spin_temp,
            halobox=halobox,
            ics=initial_conditions,
        )
    else:
        # If we have only one cell (could happen if we do a global evolution), we set the neutral fraction
        # according to the global evolution

        from .global_evolution import (
            compute_global_reionization_at_z,  # This is imported here to prevent circular import
        )

        box = compute_global_reionization_at_z(
            redshift=redshift,
            inputs=inputs,
            previous_ionized_box=previous_ionized_box,
            spin_temp=spin_temp,
        )
        return box


@single_field_func
def brightness_temperature(
    *,
    ionized_box: IonizedBox,
    perturbed_field: PerturbedField,
    spin_temp: TsBox | None = None,
) -> BrightnessTemp:
    r"""
    Compute a coeval brightness temperature box.

    Parameters
    ----------
    ionized_box: :class:`IonizedBox`
        A pre-computed ionized box.
    perturbed_field: :class:`PerturbedField`
        A pre-computed perturbed field at the same redshift as `ionized_box`.
    spin_temp: :class:`TsBox`, optional
        A pre-computed spin temperature, at the same redshift as the other boxes.

    Returns
    -------
    :class:`BrightnessTemp` instance.
    """
    redshift = ionized_box.redshift
    inputs = ionized_box.inputs

    if spin_temp is None:
        if inputs.astro_options.USE_TS_FLUCT:
            raise ValueError(
                "You have USE_TS_FLUCT=True, but have not provided a spin_temp!"
            )
        else:
            spin_temp = TsBox.dummy()

    box = BrightnessTemp.new(redshift=redshift, inputs=inputs)

    return box.compute(
        spin_temp=spin_temp,
        ionized_box=ionized_box,
        perturbed_field=perturbed_field,
    )
