class FactorEncodingError(FormulaMaterializationError):
    pass class FormulaMaterializerNotFoundError(FormulaicError):
    pass class EvaluatedFactor:
    """
    A container for the evaluated state of a `Factor` object in a given context.

    This class acts as the glue between an abstract `Factor` specification and
    the realisation of that factor in a specific data context.

    Attributes:
        factor: The `Factor` instance for which values have been computed.
        values: The evaluated values for the factor.
    """

    factor: Optional[Factor] = None
    values: Optional[FactorValues[Any]] = None

    @property
    def expr(self) -> str:
        """
        The expression of the evaluated factor.
        """
        return self.factor.expr

    @property
    def metadata(self) -> FactorValuesMetadata:
        """
        The metadata associated with the evaluated values.
        """
        return self.values.__formulaic_metadata__

    def __repr__(self) -> str:
        return repr(self.factor)

    def __eq__(self, other) -> bool:
        if isinstance(other, EvaluatedFactor):
            return self.factor == other.factor
        return NotImplemented

    def __lt__(self, other) -> bool:
        if isinstance(other, EvaluatedFactor):
            return self.factor < other.factor
        return NotImplemented

    def replace(self, **changes) -> EvaluatedFactor:
        """
        Return a copy of this `EvaluatedFactor` instance with the nominated
        attributes mutated.
        """
        return replace(self, **changes) class FactorValues(Generic[T], wrapt.ObjectProxy):
    """
    A convenience wrapper that surfaces a `FactorValuesMetadata` instance at
    `<object>.__formulaic_metadata__`. This wrapper can otherwise wrap any
    object and behaves just like that object.
    """

    def __init__(
        self,
        values,
        *,
        metadata: FactorValuesMetadata = MISSING,
        kind: Union[str, Factor.Kind] = MISSING,
        column_names: Tuple[str] = MISSING,
        spans_intercept: bool = MISSING,
        drop_field: Optional[str] = MISSING,
        format: str = MISSING,
        encoded: bool = MISSING,
    ):
        metadata_constructor = FactorValuesMetadata
        metadata_kwargs = dict(
            kind=Factor.Kind(kind) if kind is not MISSING else kind,
            column_names=column_names,
            spans_intercept=spans_intercept,
            drop_field=drop_field,
            format=format,
            encoded=encoded,
        )
        for key in set(metadata_kwargs):
            if metadata_kwargs[key] is MISSING:
                metadata_kwargs.pop(key)

        if hasattr(values, "__formulaic_metadata__"):
            metadata_constructor = values.__formulaic_metadata__.replace
            if isinstance(values, FactorValues):
                values = values.__wrapped__

        if metadata:
            metadata_constructor = metadata.replace

        wrapt.ObjectProxy.__init__(self, values)
        self._self_metadata = metadata_constructor(**metadata_kwargs)

    @property
    def __formulaic_metadata__(self) -> FactorValuesMetadata:
        return self._self_metadata

    def __repr__(self) -> str:
        return self.__wrapped__.__repr__()  # pragma: no cover class ScopedFactor:
    def __init__(self, factor, reduced=False):
        self.factor = factor
        self.reduced = reduced

    def __repr__(self):
        return repr(self.factor) + ("-" if self.reduced else "")

    def __hash__(self):
        return hash(repr(self))

    def __eq__(self, other):
        if isinstance(other, ScopedFactor):
            return self.factor == other.factor and self.reduced == other.reduced
        return NotImplemented

    def __lt__(self, other):
        if isinstance(other, ScopedFactor):
            if self.factor == other.factor:
                return self.reduced > other.reduced
            return self.factor < other.factor
        return NotImplemented class ScopedTerm:

    __slots__ = ("factors", "scale")

    def __init__(self, factors, scale=None):
        self.factors = tuple(sorted(factors))
        self.scale = scale

    def __hash__(self):
        return hash(self.factors)

    def __eq__(self, other):
        if isinstance(other, ScopedTerm):
            return self.factors == other.factors
        return NotImplemented

    def __repr__(self):
        factor_repr = (
            ":".join(f.__repr__() for f in sorted(self.factors))
            if self.factors
            else "1"
        )
        if self.scale is not None and self.scale != 1:
            return f"{self.scale}*{factor_repr}"
        return factor_repr

    def copy(self, *, without_values=False):
        factors = self.factors
        if without_values:
            factors = [
                ScopedFactor(
                    factor=factor.factor.replace(values=None),
                    reduced=factor.reduced,
                )
                for factor in factors
            ]
        return ScopedTerm(factors, scale=self.scale) class FormulaMaterializer(metaclass=FormulaMaterializerMeta):

    REGISTER_NAME = None
    REGISTER_INPUTS = set()
    REGISTER_OUTPUTS = set()
    REGISTER_PRECEDENCE = 100

    # Public API

    @inherit_docs(method="_init")
    def __init__(self, data, context=None, **kwargs):
        self.data = data
        self.context = context or {}
        self._init(**kwargs)

        self.layered_context = LayeredMapping(
            self.data_context, self.context, TRANSFORMS
        )

        self.factor_cache = {}
        self.encoded_cache = {}

    def _init(self):
        pass  # pragma: no cover

    @property
    def data_context(self):
        return self.data

    @property
    def nrows(self):
        return len(self.data)

    def get_model_matrix(
        self, spec, ensure_full_rank=True, na_action="drop", output=None
    ):
        from formulaic.formula import Formula
        from formulaic.model_spec import ModelSpec

        # Prepare ModelSpec instance
        if output is None:
            output = self.REGISTER_OUTPUTS[0]

        if isinstance(spec, Formula):
            spec = ModelSpec(
                formula=spec,
                materializer=self,
                ensure_full_rank=ensure_full_rank,
                na_action=na_action,
                output=output,
            )
        if not isinstance(spec, ModelSpec):
            spec = ModelSpec(
                formula=Formula(spec),
                materializer=self,
                ensure_full_rank=ensure_full_rank,
                na_action=na_action,
                output=output,
            )

        if output not in self.REGISTER_OUTPUTS:
            raise FormulaMaterializationError(
                f"Nominated output {repr(output)} is invalid. Available output types are: {set(self.REGISTER_OUTPUTS)}."
            )

        # Step 0: Check whether formula separators are in play, and if so, recurse.
        if isinstance(spec.formula.terms, tuple):
            return tuple(
                self.get_model_matrix(
                    Formula(terms),
                    ensure_full_rank=ensure_full_rank,
                    na_action=na_action,
                    output=output,
                )
                for terms in spec.formula.terms
            )

        drop_rows = (
            set()
        )  # Keep track of which rows to drop if `self.config.na_action == 'drop'`.

        # Step 1: Evaluate all factors
        for term in spec.formula.terms:
            for factor in term.factors:
                self._evaluate_factor(factor, spec, drop_rows)

        drop_rows = sorted(drop_rows)

        # Step 2: Determine strategy to maintain structural full-rankness of output matrix
        scoped_terms_for_terms = self._get_scoped_terms(
            spec.formula.terms, ensure_full_rank=spec.ensure_full_rank
        )

        # Step 3: Generate the columns which will be collated into the full matrix
        cols = []
        for term, scoped_terms in scoped_terms_for_terms:
            scoped_cols = OrderedDict()
            for scoped_term in scoped_terms:
                if not scoped_term.factors:
                    scoped_cols[
                        "Intercept"
                    ] = scoped_term.scale * self._encode_constant(
                        1, None, {}, spec, drop_rows
                    )
                else:
                    scoped_cols.update(
                        self._get_columns_for_term(
                            [
                                self._encode_evaled_factor(
                                    scoped_factor.factor,
                                    spec,
                                    drop_rows,
                                    reduced_rank=scoped_factor.reduced,
                                )
                                for scoped_factor in sorted(scoped_term.factors)
                            ],
                            spec=spec,
                            scale=scoped_term.scale,
                        )
                    )
            cols.append((term, scoped_terms, scoped_cols))

        # Step 4: Populate remaining model spec fields
        spec.materializer = self
        if spec.structure:
            cols = self._enforce_structure(cols, spec, drop_rows)
        else:
            spec.structure = [
                (
                    term,
                    list(st.copy(without_values=True) for st in scoped_terms),
                    list(scoped_cols),
                )
                for term, scoped_terms, scoped_cols in cols
            ]

        # Step 5: Collate factors into one ModelMatrix
        return ModelMatrix(
            self._combine_columns(
                [
                    (name, values)
                    for term, scoped_terms, scoped_cols in cols
                    for name, values in scoped_cols.items()
                ],
                spec=spec,
                drop_rows=drop_rows,
            ),
            spec=spec,
        )

    # Methods related to ensuring out matrices are structurally full-rank

    def _get_scoped_terms(self, terms, ensure_full_rank=True):
        """
        Generate the terms to be used in the model matrix.

        This method first evaluates each factor in the context of the data
        (and environment), and then determines the correct "scope" (full vs.
        reduced rank) for each term. If `ensure_full_rank` is `True`, then the
        resulting terms when combined is guaranteed to be structurally full-rank.

        Args:
            terms (list<Term>): A list of term arguments (usually from a formula)
                object.
            ensure_full_rank (bool): Whether evaluated terms should be scoped
                to ensure that their combination will result in a full-rank
                matrix.
            transform_state (dict): The state of any stateful transforms
                (will be populated if empty).

        Returns:
            list<ScopedTerm>: A list of appropriately scoped terms.
        """
        spanned = set()

        for term in terms:
            evaled_factors = [self.factor_cache[factor.expr] for factor in term.factors]

            if ensure_full_rank:
                term_span = self._get_scoped_terms_spanned_by_evaled_factors(
                    evaled_factors
                ).difference(spanned)
                scoped_terms = self._simplify_scoped_terms(term_span)
                spanned.update(term_span)
            else:
                scoped_terms = [
                    ScopedTerm(
                        factors=(
                            ScopedFactor(evaled_factor, reduced=False)
                            for evaled_factor in evaled_factors
                            if evaled_factor.metadata.kind is not Factor.Kind.CONSTANT
                        ),
                        scale=functools.reduce(
                            operator.mul,
                            [
                                evaled_factor.values
                                for evaled_factor in evaled_factors
                                if evaled_factor.metadata.kind.value
                                is Factor.Kind.CONSTANT
                            ],
                            1,
                        ),
                    )
                ]
            yield term, scoped_terms

    @classmethod
    def _get_scoped_terms_spanned_by_evaled_factors(
        cls, evaled_factors: Iterable[EvaluatedFactor]
    ) -> Set[ScopedTerm]:
        """
        Return the set of ScopedTerm instances which span the set of
        evaluated factors.

        Args:
            evaled_factors: The evaluated factors for which to generated scoped
                terms.

        Returns:
            The scoped terms for the nominated `evaled_factors`.
        """
        scale = 1
        factors = []
        for factor in evaled_factors:
            if factor.metadata.kind is Factor.Kind.CONSTANT:
                scale *= factor.values
            elif factor.metadata.spans_intercept:
                factors.append((1, ScopedFactor(factor, reduced=True)))
            else:
                factors.append((ScopedFactor(factor),))
        return set(
            ScopedTerm(factors=(p for p in prod if p != 1), scale=scale)
            for prod in itertools.product(*factors)
        )

    @classmethod
    def _simplify_scoped_terms(cls, scoped_terms):
        """
        Return the minimal set of ScopedTerm instances that spans the same vectorspace.
        """
        terms = []
        for scoped_term in sorted(scoped_terms, key=lambda x: len(x.factors)):
            factors = set(scoped_term.factors)
            combined = False
            for co_scoped_term in terms:
                cofactors = set(co_scoped_term.factors)
                factors_diff = factors.difference(cofactors)
                if len(factors) - 1 != len(cofactors) or len(factors_diff) != 1:
                    continue
                factor_new = next(iter(factors_diff))
                if factor_new.reduced:
                    co_scoped_term.factors += (
                        ScopedFactor(factor_new.factor, reduced=False),
                    )
                    terms = cls._simplify_scoped_terms(terms)
                    combined = True
                    break
            if not combined:
                terms.append(scoped_term.copy())
        return terms

    # Methods related to looking-up, evaluating and encoding terms and factors

    def _evaluate_factor(
        self, factor: Factor, spec: ModelSpec, drop_rows: set
    ) -> EvaluatedFactor:
        if factor.expr not in self.factor_cache:
            try:
                if factor.eval_method.value == "lookup":
                    value = self._lookup(factor.expr)
                elif factor.eval_method.value == "python":
                    value = self._evaluate(factor.expr, factor.metadata, spec)
                elif factor.eval_method.value == "literal":
                    value = FactorValues(
                        self._evaluate(factor.expr, factor.metadata, spec),
                        kind=Factor.Kind.CONSTANT,
                    )
                else:
                    raise FactorEvaluationError(
                        f"The evaluation method `{factor.eval_method.value}` for factor `{factor}` is not understood."
                    )
            except FactorEvaluationError:
                raise
            except Exception as e:
                raise FactorEvaluationError(
                    f"Unable to evaluate factor `{factor}`. [{type(e).__name__}: {e}]"
                )

            if not isinstance(value, FactorValues):
                value = FactorValues(value)

            if value.__formulaic_metadata__.kind is Factor.Kind.UNKNOWN:
                if self._is_categorical(value):
                    kind = Factor.Kind.CATEGORICAL
                    spans_intercept = True
                else:
                    kind = Factor.Kind.NUMERICAL
                    spans_intercept = False

                value = FactorValues(value, kind=kind, spans_intercept=spans_intercept)

            if (
                factor.kind is not Factor.Kind.UNKNOWN
                and factor.kind is not value.__formulaic_metadata__.kind
            ):
                if factor.kind is Factor.Kind.CATEGORICAL:
                    value.__formulaic_metadata__.kind = factor.kind
                else:
                    raise FactorEncodingError(
                        f"Factor `{factor}` is expecting values of kind '{factor.kind.value}', "
                        f"but they are actually of kind '{value.__formulaic_metadata__.kind.value}'."
                    )
            if (
                factor.expr in spec.encoder_state
                and value.__formulaic_metadata__.kind
                is not spec.encoder_state[factor.expr][0]
            ):
                raise FactorEncodingError(
                    f"The model specification expects factor `{factor}` to have values of kind "
                    f"`{spec.encoder_state[factor.expr][0]}`, but they are actually of kind "
                    f"`{value.__formulaic_metadata__.kind.value}`."
                )
            self._check_for_nulls(factor.expr, value, spec.na_action, drop_rows)
            self.factor_cache[factor.expr] = EvaluatedFactor(
                factor=factor, values=value
            )
        return self.factor_cache[factor.expr]

    def _lookup(self, name):
        return self.layered_context[name]

    def _evaluate(self, expr, metadata, spec):
        return stateful_eval(
            expr, self.layered_context, {expr: metadata}, spec.transform_state, spec
        )

    def _is_categorical(self, values):
        if hasattr(values, "__formulaic_metadata__"):
            return values.__formulaic_metadata__.kind is Factor.Kind.CATEGORICAL
        return False

    def _check_for_nulls(self, name, values, na_action, drop_rows):
        pass  # pragma: no cover

    def _encode_evaled_factor(
        self,
        factor: EvaluatedFactor,
        spec: ModelSpec,
        drop_rows: set,
        reduced_rank: bool = False,
    ) -> Dict[str, Any]:
        if not isinstance(factor.values, dict) or not factor.metadata.encoded:
            if factor.expr in self.encoded_cache:
                encoded = self.encoded_cache[factor.expr]
            elif (factor.expr, reduced_rank) in self.encoded_cache:
                encoded = self.encoded_cache[(factor.expr, reduced_rank)]
            else:

                def map_dict(f):
                    """
                    This decorator allows an encoding function to operator on
                    dictionaries (which should be mapped over). This allows
                    transforms to output multiple non-encoded columns and still
                    have everything work as expected.
                    """

                    @functools.wraps(f)
                    def wrapped(values, metadata, state, *args, **kwargs):
                        if isinstance(values, dict):
                            encoded = {}
                            for k, v in values.items():
                                if isinstance(k, str) and k.startswith("__"):
                                    encoded[k] = v
                                else:
                                    nested_state = state.get(k, {})
                                    encoded[k] = wrapped(
                                        v, metadata, nested_state, *args, **kwargs
                                    )
                                    if nested_state:
                                        state[k] = nested_state
                            return encoded
                        return f(values, metadata, state, *args, **kwargs)

                    return wrapped

                # If we need to unpack values into columns, we do this here.
                # Otherwise, we pass through the original values.
                factor_values = FactorValues(
                    self._extract_columns_for_encoding(factor),
                    metadata=factor.metadata,
                )

                encoder_state = spec.encoder_state.get(factor.expr, [None, {}])[1]
                if factor.metadata.kind is Factor.Kind.CATEGORICAL:
                    encoded = map_dict(self._encode_categorical)(
                        factor_values,
                        factor.metadata,
                        encoder_state,
                        spec,
                        drop_rows,
                        reduced_rank=reduced_rank,
                    )
                elif factor.metadata.kind is Factor.Kind.NUMERICAL:
                    encoded = map_dict(self._encode_numerical)(
                        factor_values, factor.metadata, encoder_state, spec, drop_rows
                    )
                elif factor.metadata.kind is Factor.Kind.CONSTANT:
                    encoded = map_dict(self._encode_constant)(
                        factor_values, factor.metadata, encoder_state, spec, drop_rows
                    )
                else:
                    raise FactorEncodingError(
                        factor
                    )  # pragma: no cover; it is not currently possible to reach this sentinel
                spec.encoder_state[factor.expr] = (factor.metadata.kind, encoder_state)

                # Only encode once for encodings where we can just drop a field
                # later on below.
                if isinstance(encoded, dict) and factor.metadata.drop_field:
                    cache_key = factor.expr
                else:
                    cache_key = (factor.expr, reduced_rank)

                self.encoded_cache[cache_key] = encoded
        else:
            encoded = factor.values

        encoded = FactorValues(
            encoded,
            metadata=getattr(encoded, "__formulaic_metadata__", factor.metadata),
            encoded=True,
        )

        # Encoded factors will now all be dicts
        if (
            isinstance(encoded, dict)
            and encoded.__formulaic_metadata__.spans_intercept
            and reduced_rank
        ):
            encoded = FactorValues(
                encoded.copy(), metadata=encoded.__formulaic_metadata__
            )
            del encoded[encoded.__formulaic_metadata__.drop_field]

        return self._flatten_encoded_evaled_factor(factor.expr, encoded)

    def _extract_columns_for_encoding(
        self, factor: EvaluatedFactor
    ) -> Union[Any, Dict[str, Any]]:
        """
        If incoming factor has values that need to be unpacked into columns
        (e.g. a two-dimensions numpy array), do that expansion here. Otherwise,
        return the current factor values.
        """
        return as_columns(factor.values)

    def _flatten_encoded_evaled_factor(
        self, name: str, values: FactorValues[dict]
    ) -> Dict[str, Any]:
        if not isinstance(values, dict):
            return {name: values}

        # Some nested dictionaries may not be a `FactorValues[dict]` instance,
        # in which case we impute the default formatter in `FactorValues.format`.
        if hasattr(values, "__formulaic_metadata__"):
            name_format = values.__formulaic_metadata__.format
        else:
            name_format = FactorValuesMetadata.format

        flattened = {}
        for subfield, value in values.items():
            if isinstance(subfield, str) and subfield.startswith("__"):
                continue
            subname = name_format.format(name=name, field=subfield)
            if isinstance(value, dict):
                flattened.update(self._flatten_encoded_evaled_factor(subname, value))
            else:
                flattened[subname] = value

        return flattened

    @abstractmethod
    def _encode_constant(self, value, metadata, encoder_state, spec, drop_rows):
        pass  # pragma: no cover

    @abstractmethod
    def _encode_categorical(
        self, values, metadata, encoder_state, spec, drop_rows, reduced_rank=False
    ):
        pass  # pragma: no cover

    @abstractmethod
    def _encode_numerical(self, values, metadata, encoder_state, spec, drop_rows):
        pass  # pragma: no cover

    # Methods related to ModelMatrix output

    def _enforce_structure(
        self,
        cols: List[Tuple[Term, List[ScopedTerm], Dict[str, Any]]],
        spec,
        drop_rows: set,
    ) -> Generator[Tuple[Term, List[ScopedTerm], Dict[str, Any]]]:
        # TODO: Verify that imputation strategies are intuitive and make sense.
        assert len(cols) == len(spec.structure)
        for i in range(len(cols)):
            scoped_cols = cols[i][2]
            target_cols = spec.structure[i][2]
            if len(scoped_cols) > len(target_cols):
                raise FactorEncodingError(
                    f"Term `{cols[i][0]}` has generated too many columns compared to specification: generated {list(scoped_cols)}, expecting {target_cols}."
                )
            elif len(scoped_cols) < len(target_cols):
                if len(scoped_cols) == 0:
                    col = self._encode_constant(0, None, None, spec, drop_rows)
                elif len(scoped_cols) == 1:
                    col = next(iter(scoped_cols.values()))
                else:
                    raise FactorEncodingError(
                        f"Term `{cols[i][0]}` has generated insufficient columns compared to specification: generated {list(scoped_cols)}, expecting {target_cols}."
                    )
                scoped_cols = {name: col for name in target_cols}
            elif set(scoped_cols) != set(target_cols):
                raise FactorEncodingError(
                    f"Term `{cols[i][0]}` has generated columns that are inconsistent with specification: generated {list(scoped_cols)}, expecting {target_cols}."
                )

            yield cols[i][0], cols[i][1], {col: scoped_cols[col] for col in target_cols}

    def _get_columns_for_term(self, factors, spec, scale=1):
        """
        Assemble the columns for a model matrix given factors and a scale.

        This performs the row-wise Kronecker product of the factors.

        Args:
            factors
            scale

        Returns:
            dict
        """
        out = OrderedDict()
        for product in itertools.product(*(factor.items() for factor in factors)):
            out[":".join(p[0] for p in product)] = scale * functools.reduce(
                operator.mul, (p[1] for p in product)
            )
        return out

    @abstractmethod
    def _combine_columns(self, cols, spec, drop_rows):
        pass  # pragma: no cover class PandasMaterializer(FormulaMaterializer):

    REGISTER_NAME = "pandas"
    REGISTER_INPUTS = ("pandas.core.frame.DataFrame",)
    REGISTER_OUTPUTS = ("pandas", "numpy", "sparse")

    @override
    def _is_categorical(self, values):
        if isinstance(values, (pandas.Series, pandas.Categorical)):
            return values.dtype == object or isinstance(
                values.dtype, pandas.CategoricalDtype
            )
        return super()._is_categorical(values)

    @override
    def _check_for_nulls(self, name, values, na_action, drop_rows):

        if na_action is NAAction.IGNORE:
            return

        if isinstance(values, dict):
            for key, vs in values.items():
                self._check_for_nulls(f"{name}[{key}]", vs, na_action, drop_rows)

        elif na_action is NAAction.RAISE:
            if isinstance(values, pandas.Series) and values.isnull().values.any():
                raise ValueError(f"`{name}` contains null values after evaluation.")

        elif na_action is NAAction.DROP:
            if isinstance(values, pandas.Series):
                drop_rows.update(numpy.flatnonzero(values.isnull().values))

        else:
            raise ValueError(
                f"Do not know how to interpret `na_action` = {repr(na_action)}."
            )  # pragma: no cover; this is currently impossible to reach

    @override
    def _encode_constant(self, value, metadata, encoder_state, spec, drop_rows):
        if spec.output == "sparse":
            return spsparse.csc_matrix(
                numpy.array([value] * self.nrows).reshape(
                    (self.nrows - len(drop_rows), 1)
                )
            )
        series = value * numpy.ones(self.nrows - len(drop_rows))
        return series

    @override
    def _encode_numerical(self, values, metadata, encoder_state, spec, drop_rows):
        if drop_rows:
            values = values.drop(index=values.index[drop_rows])
        if spec.output == "sparse":
            return spsparse.csc_matrix(numpy.array(values).reshape((self.nrows, 1)))
        return values

    @override
    def _encode_categorical(
        self, values, metadata, encoder_state, spec, drop_rows, reduced_rank=False
    ):
        # Even though we could reduce rank here, we do not, so that the same
        # encoding can be cached for both reduced and unreduced rank. The
        # rank will be reduced in the _encode_evaled_factor method.
        from formulaic.transforms import encode_categorical

        if drop_rows:
            values = values.drop(index=values.index[drop_rows])
        return encode_categorical(
            values,
            reduced_rank=False,
            _metadata=metadata,
            _state=encoder_state,
            _spec=spec,
        )

    @override
    def _get_columns_for_term(self, factors, spec, scale=1):
        out = OrderedDict()

        # Pre-multiply factors with only one set of values (improves performance)
        solo_factors = {}
        indices = []
        for i, factor in enumerate(factors):
            if len(factor) == 1:
                solo_factors.update(factor)
                indices.append(i)
        if solo_factors:
            for index in reversed(indices):
                factors.pop(index)
            if spec.output == "sparse":
                factors.append(
                    {
                        ":".join(solo_factors): functools.reduce(
                            spsparse.csc_matrix.multiply, solo_factors.values()
                        )
                    }
                )
            else:
                factors.append(
                    {
                        ":".join(solo_factors): functools.reduce(
                            lambda a, b: numpy.multiply(a, b),
                            (p for p in solo_factors.values()),
                        )
                    }
                )

        for product in itertools.product(*(factor.items() for factor in factors)):
            if spec.output == "sparse":
                out[":".join(p[0] for p in product)] = scale * functools.reduce(
                    spsparse.csc_matrix.multiply, (p[1] for p in product)
                )
            else:
                out[":".join(p[0] for p in product)] = scale * functools.reduce(
                    lambda a, b: numpy.multiply(a, b),
                    (numpy.array(p[1]) for p in product),
                )
        return out

    @override
    def _combine_columns(self, cols, spec, drop_rows):

        # If we are outputing a pandas DataFrame, explicitly override index
        # in case transforms/etc have lost track of it.
        if spec.output == "pandas":
            pandas_index = self.data_context.index
            if drop_rows:
                pandas_index = pandas_index.drop(self.data_context.index[drop_rows])

        # Special case no columns to empty csc_matrix, array, or DataFrame
        if not cols:
            values = numpy.empty((self.data.shape[0], 0))
            if spec.output == "sparse":
                return spsparse.csc_matrix(values)
            elif spec.output == "numpy":
                return values
            else:
                return pandas.DataFrame(index=pandas_index)

        # Otherwise, concatenate columns into model matrix
        if spec.output == "sparse":
            return spsparse.hstack([col[1] for col in cols])
        if spec.output == "numpy":
            return numpy.stack([col[1] for col in cols], axis=1)
        return pandas.DataFrame(
            {col[0]: col[1] for col in cols},
            index=pandas_index,
            copy=False,
        ) class ModelSpec:
    def __init__(
        self,
        formula,
        ensure_full_rank=True,
        structure=None,
        materializer=None,
        na_action="drop",
        output=None,
        transform_state=None,
        encoder_state=None,
    ):
        self.formula = Formula.from_spec(formula)
        self.ensure_full_rank = ensure_full_rank
        self.structure = structure
        self.materializer = materializer
        self.na_action = NAAction(na_action)
        self.output = output
        self.transform_state = transform_state or {}
        self.encoder_state = encoder_state or {}

    @property
    def materializer(self):
        return self._materializer

    @materializer.setter
    def materializer(self, materializer):
        if (
            isinstance(materializer, FormulaMaterializer)
            or inspect.isclass(materializer)
            and issubclass(materializer, FormulaMaterializer)
        ):
            materializer = materializer.REGISTER_NAME
        assert materializer is None or isinstance(materializer, str), materializer
        self._materializer = materializer

    @property
    def feature_names(self):
        return [name for row in self.structure for name in row[2]]

    @property
    def feature_indices(self):
        return OrderedDict([(name, i) for i, name in enumerate(self.feature_names)])

    @property
    def term_slices(self):
        slices = OrderedDict()
        start = 0
        for row in self.structure:
            end = start + len(row[2])
            slices[row[0]] = slice(start, end)
            start = end
        return slices

    def get_model_matrix(self, data, **kwargs):
        if self.materializer is None:
            materializer = FormulaMaterializer.for_data(data)
        else:
            materializer = FormulaMaterializer.for_materializer(self.materializer)
        return materializer(data, **kwargs).get_model_matrix(self)

    def differentiate(self, *vars, use_sympy=False):
        return ModelSpec(
            formula=self.formula.differentiate(*vars, use_sympy=use_sympy),
            ensure_full_rank=self.ensure_full_rank,
            structure=self.structure,
            materializer=self.materializer,
            transform_state=self.transform_state,
            encoder_state=self.encoder_state,
        ) class Factor:
    """
    Factors are the indivisable atomic unit that make up formulas.

    Each instance of `Factor` is a specification that is evaluable by a
    materializer to generate concrete vector(s). `Factors` are multiplied
    together into `Term`s, which in turn represent the output columns of model
    matrices. Note that `Factor` instances are entirely abstract of data.

    Attributes:
        expr: The (string) expression to be evaluated by the materializer.
        eval_method: An `EvalMethod` enum instance indicating the mechanism to
            be used to evaluate the expression (one of: unknown, literal, lookup
            or python).
        kind: The kind of data represented (one of: unknown, constant,
            numerical, categorical).
        metadata: An additional (optional) dictionary of metadata (currently
            unused).
    """

    class EvalMethod(Enum):
        UNKNOWN = "unknown"
        LITERAL = "literal"
        LOOKUP = "lookup"
        PYTHON = "python"

    class Kind(Enum):
        UNKNOWN = "unknown"
        CONSTANT = "constant"
        NUMERICAL = "numerical"
        CATEGORICAL = "categorical"

    __slots__ = ("expr", "_eval_method", "_kind", "metadata")

    def __init__(
        self,
        expr: str = "",
        *,
        eval_method: Optional[Union[str, EvalMethod]] = None,
        kind: Optional[Union[str, Kind]] = None,
        metadata: Optional[Dict] = None
    ):
        self.expr = expr
        self.eval_method = eval_method
        self.kind = kind
        self.metadata = metadata or {}

    @property
    def eval_method(self) -> EvalMethod:
        return self._eval_method

    @eval_method.setter
    def eval_method(self, eval_method):
        self._eval_method = Factor.EvalMethod(eval_method or "unknown")

    @property
    def kind(self) -> Kind:
        return self._kind

    @kind.setter
    def kind(self, kind):
        self._kind = Factor.Kind(kind or "unknown")

    def __eq__(self, other):
        if isinstance(other, str):
            return self.expr == other
        if isinstance(other, Factor):
            return self.expr == other.expr
        return NotImplemented

    def __hash__(self):
        return self.expr.__hash__()

    def __lt__(self, other):
        if isinstance(other, Factor):
            return self.expr < other.expr
        return NotImplemented

    def to_terms(self) -> Iterable[Term]:
        """
        Convert this `Factor` instance into a `Term` instance, and expose it as
        a single-element iterable.
        """
        return {Term([self])}

    def __repr__(self):
        return self.expr
import pandas
import pytest
from formulaic.errors import FactorEncodingError, FormulaMaterializerNotFoundError
from formulaic.materializers.types import (
    EvaluatedFactor,
    FactorValues,
    ScopedFactor,
    ScopedTerm,
)
from formulaic.materializers.base import FormulaMaterializer
from formulaic.materializers.pandas import PandasMaterializer
from formulaic.model_spec import ModelSpec
from formulaic.parser.types import Factor



class TestFormulaMaterializer:
    def test_registrations(self):
        assert sorted(FormulaMaterializer.REGISTERED_NAMES) == ["arrow", "pandas"]
        assert sorted(FormulaMaterializer.REGISTERED_INPUTS) == [
            "pandas.core.frame.DataFrame",
            "pyarrow.lib.Table",
        ]

    def test_retrieval(self):
        assert FormulaMaterializer.for_materializer("pandas") is PandasMaterializer
        assert (
            FormulaMaterializer.for_materializer(PandasMaterializer)
            is PandasMaterializer
        )
        assert (
            FormulaMaterializer.for_data(pandas.DataFrame(), output="numpy")
            is PandasMaterializer
        )

        with pytest.raises(FormulaMaterializerNotFoundError):
            FormulaMaterializer.for_materializer("invalid_materializer")

        with pytest.raises(
            FormulaMaterializerNotFoundError,
            match=r"No materializer has been registered for input type",
        ):
            FormulaMaterializer.for_data("str")

        with pytest.raises(
            FormulaMaterializerNotFoundError,
            match=r"No materializer has been registered for input type .* that supports output type",
        ):
            FormulaMaterializer.for_data(pandas.DataFrame(), output="invalid_output")

    def test_defaults(self):
        assert (

            FormulaMaterializer.for_data(pandas.DataFrame({"a": [1, 2, 3]}))