def construct_default_space(mf: SCF) -> Union[RConstructSpaceReturnType, UConstructSpaceReturnType]:
    """Construct a default space.

    Args:
        mf: PySCF mean-field object.

    Returns:
        The molecular orbital coefficients, the molecular orbital occupation numbers, and the
        default space.
    """

    def _construct(mo_occ: NDArray[T]) -> Space:
        """Build the default space."""
        frozen = np.zeros(mo_occ.shape, dtype=bool)
        active = np.zeros(mo_occ.shape, dtype=bool)
        space = Space(
            occupied=mo_occ > 0,
            frozen=frozen,
            active=active,
        )
        return space

    # Construct the default space
    if mf.mo_occ.ndim == 2:
        space_a = _construct(mf.mo_occ[0])
        space_b = _construct(mf.mo_occ[1])
        return mf.mo_coeff, mf.mo_occ, (space_a, space_b)
    else:
        return mf.mo_coeff, mf.mo_occ, _construct(mf.mo_occ)

def construct_fno_space(
    mf: SCF,
    occ_tol: Optional[float] = 1e-5,
    occ_frac: Optional[float] = None,
    amplitudes: Optional[Namespace[SpinArrayType]] = None,
) -> Union[RConstructSpaceReturnType, UConstructSpaceReturnType]:
    """Construct a frozen natural orbital space.

    Args:
        mf: PySCF mean-field object.
        occ_tol: Threshold in the natural orbital occupation numbers.
        occ_frac: Fraction of the natural orbital occupation numbers to be retained. Overrides
            `occ_tol` if both are specified.
        amplitudes: Cluster amplitudes. If provided, use these amplitudes when calculating the MP2
            1RDM.

    Returns:
        The natural orbital coefficients, the natural orbital occupation numbers, and the frozen
        natural orbital space.
    """
    # Get the MP2 1RDM
    solver = MP2(mf)
    if not amplitudes:
        solver.kernel()
        dm1 = solver.make_rdm1().astype(types[float])
    else:
        if isinstance(amplitudes.t2, util.Namespace):
            t2 = (amplitudes.t2.aaaa, amplitudes.t2.abab, amplitudes.t2.bbbb)
            dm1 = solver.make_rdm1(t2=t2).astype(types[float])
        else:
            dm1 = solver.make_rdm1(t2=amplitudes.t2).astype(types[float])

    # def _construct(dm1, mo_energy, mo_coeff, mo_occ):
    def _construct(
        dm1: NDArray[T],
        mo_energy: NDArray[T],
        mo_coeff: NDArray[T],
        mo_occ: NDArray[T],
    ) -> RConstructSpaceReturnType:
        # Get the number of occupied orbitals
        nocc = (mo_occ > 0).astype(int).sum()

        # Calculate the natural orbitals
        n, c = np.linalg.eigh(dm1[nocc:, nocc:])
        n, c = n[::-1], c[:, ::-1]

        # Truncate the natural orbitals
        if occ_frac is None:
            active_vir = n > occ_tol
        else:
            active_vir = np.cumsum(n / n.sum()) <= occ_frac
        num_active_vir = active_vir.astype(int).sum()

        # Canonicalise the natural orbitals
        fock_vv = np.diag(mo_energy[nocc:])
        fock_vv = util.einsum("ab,au,bv->uv", fock_vv, c, c)
        _, c_can = np.linalg.eigh(fock_vv[active_vir][:, active_vir])

        # Transform the natural orbitals
        no_coeff_avir = util.einsum(
            "pi,iq,qj->pj", mo_coeff[:, nocc:], c[:, :num_active_vir], c_can
        )
        no_coeff_fvir = mo_coeff[:, nocc:] @ c[:, num_active_vir:]
        no_coeff_occ = mo_coeff[:, :nocc]
        no_coeff = np.concatenate((no_coeff_occ, no_coeff_avir, no_coeff_fvir), axis=1)

        # Build the natural orbital space
        active = np.zeros(mo_occ.shape, dtype=bool)
        frozen = np.zeros(mo_occ.shape, dtype=bool)
        frozen[nocc + num_active_vir :] = True
        no_space = Space(
            occupied=mo_occ > 0,
            frozen=frozen,
            active=active,
        )

        return no_coeff, mo_occ, no_space

    # Construct the natural orbitals
    if mf.mo_occ.ndim == 2:
        coeff_a, occ_a, space_a = _construct(
            dm1[0].astype(types[float]),
            mf.mo_energy[0].astype(types[float]),
            mf.mo_coeff[0].astype(types[float]),
            mf.mo_occ[0].astype(types[float]),
        )
        coeff_b, occ_b, space_b = _construct(
            dm1[1].astype(types[float]),
            mf.mo_energy[1].astype(types[float]),
            mf.mo_coeff[1].astype(types[float]),
            mf.mo_occ[1].astype(types[float]),
        )
        return (coeff_a, coeff_b), (occ_a, occ_b), (space_a, space_b)
    else:
        return _construct(
            dm1.astype(types[float]),
            mf.mo_energy.astype(types[float]),
            mf.mo_coeff.astype(types[float]),
            mf.mo_occ.astype(types[float]),
        )