import warnings
import copy
from collections.abc import Callable
import numpy as np
import scipy as scp
import pandas as pd
from tqdm import tqdm
from .smooths import TP_basis_calc
from .terms import GammTerm, i, l, f, irf, ri, rs, fs
from .penalties import embed_in_S_sparse, embed_in_Sj_sparse
from .file_loading import (
    read_cov,
    read_cor_cov_single,
    read_cov_no_cor,
    read_unique,
    read_dtype,
    mp,
    repeat,
)
from .custom_types import (
    PenType,
    LambdaTerm,
    ConstType,
    Constraint,
    VarType,
)
from .matrix_solvers import translate_sparse, eigen_solvers


class lhs:
    """
    The Left-hand side of a regression equation.

    See the :class:`Formula` class for examples.

    :param variable: The name of the dependent/response variable in the dataframe passed to a
        :class:`Formula`. Can point to continuous and categorical variables. For
        :class:`mssm..models.GSMM` models, the variable can also be set to any placeholder
        variable in the data, since not every :class:`Formula` will be associated with a
        particular response variable.
    :type variable: str
    :param f: A function that will be applied to the ``variable`` before fitting. For example:
        np.log(). By default no function is applied to the ``variable``.
    :type f: Callable, optional
    """

    def __init__(self, variable: str, f: Callable = None) -> None:
        self.variable = variable
        self.f = f


def _compute_constraint_single_MP(
    sterm: f,
    vars: list[str],
    lhs_var: str,
    file: str,
    var_mins: dict,
    var_maxs: dict,
    file_loading_kwargs: dict,
) -> np.ndarray:
    """Internal function to compute QR identifiability constraint based on reading model matrix
    data from file

    Wood (2017) has an overview over identifiability constraintsfor smooth terms.

    References:
       - Wood, S. N. (2017). Generalized Additive Models: An Introduction with R, Second Edition \
        (2nd ed.).

    :param sterm: Smooth term to make identifiable
    :type sterm: f
    :param vars: List of variables
    :type vars: list[str]
    :param lhs_var: name of dependent variable
    :type lhs_var: str
    :param file: file name
    :type file: str
    :param var_mins: Dictionary holding covariate minimums
    :type var_mins: dict
    :param var_maxs: Dictionary holding covariate maximums
    :type var_maxs: dict
    :param file_loading_kwargs: Any optional file loading key-word arguments.
    :type file_loading_kwargs: dict
    :return: Constraint vector ``C``
    :rtype: np.ndarray
    """

    C = 0

    if len(vars) > 1 and sterm.te is False:
        C = [0 for _ in range(len(vars))]

    matrix_term = None  # for Te basis
    for vi in range(len(vars)):
        # If a smooth term needs to be identifiable I act as if you would have asked for nk+1!
        # so that the identifiable term is of the dimension expected.

        if len(vars) > 1:
            id_nk = sterm.nk[vi]
        else:
            id_nk = sterm.nk

        if sterm.te is False:
            id_nk += 1

        var_cov_flat = read_cor_cov_single(lhs_var, vars[vi], file, file_loading_kwargs)

        matrix_term_v = sterm.basis(
            var_cov_flat,
            None,
            id_nk,
            min_c=var_mins[vars[vi]],
            max_c=var_maxs[vars[vi]],
            **sterm.basis_kwargs,
        )

        if sterm.te is False:
            if len(vars) > 1:
                C[vi] += np.sum(matrix_term_v, axis=0).reshape(-1, 1)
            else:
                C += np.sum(matrix_term_v, axis=0).reshape(-1, 1)
        else:
            if vi == 0:
                matrix_term = matrix_term_v
            else:
                matrix_term = TP_basis_calc(matrix_term, matrix_term_v)

    # Now deal with te basis
    if sterm.te:
        C += np.sum(matrix_term, axis=0).reshape(-1, 1)

    return C


class Formula:
    """The formula of a regression equation.

    **Note:** The class implements multiple ``get_*`` functions to access attributes stored in
    instance variables. The get functions always return a copy of the instance variable and the
    results are thus safe to manipulate.

    Examples::

      from mssm.models import *
      from mssmViz.sim import *

      from mssm.src.python.formula import build_penalties,build_model_matrix

      # Get some data and formula
      Binomdat = sim3(10000,0.1,family=Binomial(),seed=20)
      formula = Formula(lhs("y"),[i(),f(["x0"]),f(["x1"]),f(["x2"]),f(["x3"])],data=Binomdat)

      # Now with a tensor smooth
      formula = Formula(lhs("y"),[i(),f(["x0","x1"],te=True),f(["x2"]),f(["x3"])],data=Binomdat)

      # Now with a tensor smooth anova style
      formula = Formula(lhs("y"),[i(),f(["x0"]),f(["x1"]),f(["x0","x1"]),f(["x2"]),f(["x3"])],
        data=Binomdat)


      ######## Stream data from file and set up custom codebook #########

      file_paths = [
        f'https://raw.githubusercontent.com/JoKra1/mssmViz/main/data/GAMM/sim_dat_cond_{cond}.csv' \
        for cond in ["a","b"]
        ]

      # Set up specific coding for factor 'cond'
      codebook = {'cond':{'a': 0, 'b': 1}}

      formula = Formula(lhs=lhs("y"), # The dependent variable - here y!
                        terms=[i(), # The intercept, a
                                 l(["cond"]), # For cond='b'
                                 # to-way interaction between time and cond;
                                 # one smooth over time per cond level
                                 f(["time"],by="cond"),
                                 # to-way interaction between x and cond;
                                 # one smooth over x per cond level
                                 f(["x"],by="cond"),
                                 # three-way interaction
                                 f(["time","x"],by="cond"),
                                 # Random non-linear effect of time - one
                                 # smooth per level of factor sub
                                 fs(["time"],rf="sub")],
                        data=None, # No data frame!
                        file_paths=file_paths, # Just a list with paths to files.
                        print_warn=False,
                        codebook=codebook)

      # Alternative:
      formula = Formula(lhs=lhs("y"),
                              terms=[i(),
                                    l(["cond"]),
                                    f(["time"],by="cond"),
                                    f(["x"],by="cond"),
                                    f(["time","x"],by="cond"),
                                    fs(["time"],rf="sub")],
                              data=None,
                              file_paths=file_paths,
                              print_warn=False,
                              keep_cov=True, # Keep encoded data structure in memory
                              codebook=codebook)

      # preparing for ar1 model (with resets per time-series) and data type requirements

      dat = pd.read_csv(
        'https://raw.githubusercontent.com/JoKra1/mssmViz/main/data/GAMM/sim_dat.csv'
        )

      # mssm requires that the data-type for variables used as factors is 'O'=object
      dat = dat.astype({'series': 'O',
                        'cond':'O',
                        'sub':'O',
                        'series':'O'})

      formula = Formula(lhs=lhs("y"),
                        terms=[i(),
                                 l(["cond"]),
                                 f(["time"],by="cond"),
                                 f(["x"],by="cond"),
                                 f(["time","x"],by="cond")],
                        data=dat,
                        print_warn=False,
                        series_id='series') # 'series' variable identifies individual time-series

    :param lhs: The lhs object defining the dependent variable.
    :type variable: lhs
    :param terms: A list of the terms which should be added to the model.
        See :py:mod:`mssm.src.python.terms` for info on which terms can be added.
    :type terms: [GammTerm]
    :param data: A pandas dataframe (with header!) of the data which should be used to estimate the
        model. The variable specified for ``lhs`` as well as all variables included for a ``term``
        in ``terms`` need to be present in the data, otherwise the call to Formula will throw an
        error.
    :type data: pd.DataFrame or None
    :param series_id: A string identifying the individual experimental units. Usually a unique
        trial identifier. Only necessary if approximate derivative computations are to be utilized
        for random smooth terms or if you need to estimate an 'ar1' model for multiple time-series
        data.
    :type series_id: str, optional
    :param codebook: Codebook - keys should correspond to factor variable names specified in terms.
        Values should again be a ``dict``, with keys for each of K levels of the factor and value
        corresponding to an integer in {0,K}.
    :type codebook: dict or None
    :param print_warn: Whether warnings should be printed. Useful when fitting models from terminal.
        Defaults to True.
    :type print_warn: bool,optional
    :param keep_cov: Whether or not the internal encoding structure of all predictor variables
        should be created when forming :math:`\\mathbf{X}^T\\mathbf{X}` iteratively instead of
        forming :math:`\\mathbf{X}` directly. Can speed up estimation but increases memory
        footprint. Defaults to True.
    :type keep_cov: bool,optional
    :param find_nested: Whether or not to check for nested smooth terms. This only has an effect if
        you include at least one smooth term with more than two variables. Additionally, this check
        is often not necessary if you correctly use the ``te`` key-word of smooth terms and ensure
        that the marginals used to construct ti smooth terms have far fewer basis functions than
        the "main effect" univariate smooths. Thus, if you know what you're doing and you're
        working with large models, you might want to disable this (i.e., set to False) because this
        check can get quite expensive for larger models. Defaults to True.
    :type find_nested: bool,optional
    :param file_paths: A list of paths to .csv files from which :math:`\\mathbf{X}^T\\mathbf{X}`
        and :math:`\\mathbf{X}^T\\mathbf{y}` should be created iteratively. Setting this to a
        non-empty list will prevent fitting X as a whole. ``data`` should then be set to ``None``.
        Defaults to an empty list.
    :type file_paths: [str],optional
    :param file_loading_nc: How many cores to use to a) accumulate :math:`\\mathbf{X}` in parallel
        (if ``data`` is not ``None`` and ``file_paths`` is an empty list) or b) to accumulate
        :math:`\\mathbf{X}^T\\mathbf{X}` and :math:`\\mathbf{X}^T\\mathbf{y}`
        (and :math:`\\mathbf{\\eta}` during estimation) (if ``data`` is ``None`` and ``file_paths``
        is a non-empty list). For case b, this should really be set to the maimum number of cores
        available. For a this only really speeds up accumulating :math:`\\mathbf{X}` if
        :math:`\\mathbf{X}` has many many columns and/or rows. Defaults to 1.
    :type file_loading_nc: int,optional
    :param file_loading_kwargs: Any key-word arguments to pass to pandas.read_csv when
        :math:`\\mathbf{X}^T\\mathbf{X}` and :math:`\\mathbf{X}^T\\mathbf{y}` should be created
        iteratively (if ``data`` is ``None`` and ``file_paths`` is a non-empty list).
        Defaults to ``{"header":0,"index_col":False}``.
    :type file_loading_kwargs: dict,optional
    :ivar lhs lhs: The left-hand side object of the regression formula passed to the constructor.
        Initialized at construction.
    :ivar [GammTerm] terms: The list of terms passed to the constructor. Initialized at
        construction.
    :ivar pd.DataFrame data: The dataframe passed to the constructor. Initialized at
        construction.
    :ivar [int] coef_per_term: A list containing the number of coefficients corresponding to each
        term included in ``terms``. Initialized at construction.
    :ivar [str] coef_names: A list containing a named identifier (e.g., "Intercept") for each
        coefficient estimated by the model. Initialized at construction.
    :ivar int n_coef: The number of coefficients estimated by the model in total. Initialized at
        construction.
    :ivar int unpenalized_coef: The number of un-penalized coefficients estimated by the model.
        Initialized at construction.
    :ivar np.ndarray or None y_flat: An array, containing all values on the dependent variable
        (i.e., specified by ``lhs.variable``) in order of the data-frame passed to ``data``. This
        variable will be initialized at construction but only if ``file_paths=None``, i.e., in case
        :math:`\\mathbf{X}^T\\mathbf{X}` and :math:`\\mathbf{X}^T\\mathbf{y}` are **not** created
        iteratively.
    :ivar np.ndarray or None cov_flat: An array, containing all (encoded, in case of categorical
        predictors) values on each predictor (each columns of ``cov_flat`` corresponds to a
        different predictor) variable included in any of the ``terms`` in order of the data-frame
        passed to ``data``. This variable will be initialized at construction but only if
        ``file_paths=None``, i.e., in case :math:`\\mathbf{X}^T\\mathbf{X}` and
        :math:`\\mathbf{X}^T\\mathbf{y}` are **not** created iteratively.
    :ivar np.ndarray or None NOT_NA_flat: An array, containing an indication (as bool) for each
        value on the dependent variable (i.e., specified by ``lhs.variable``) whether the
        corresponding value is not a number ("NA") or not. In order of the data-frame passed to
        ``data``. This variable will be initialized at construction but only if ``file_paths=None``,
        i.e., in case :math:`\\mathbf{X}^T\\mathbf{X}` and :math:`\\mathbf{X}^T\\mathbf{y}` are
        **not** created iteratively.
    """

    def __init__(
        self,
        lhs: lhs,
        terms: list[GammTerm],
        data: pd.DataFrame,
        series_id: str | None = None,
        codebook: dict | None = None,
        print_warn: bool = True,
        keep_cov: bool = False,
        find_nested: bool = True,
        file_paths: list[str] = [],
        file_loading_nc: int = 1,
        file_loading_kwargs: dict = {"header": 0, "index_col": False},
    ) -> None:

        self.lhs = lhs
        self.terms = terms
        self.data = data
        self.series_id = series_id
        self.print_warn = print_warn
        self.keep_cov = keep_cov  # For iterative X.T@X building
        self.file_paths = (
            file_paths  # If this will not be empty, we accumulate t(X)@X directly
        )
        self.file_loading_nc = file_loading_nc
        self.file_loading_kwargs = file_loading_kwargs
        self.factor_codings = {}
        self.coding_factors = {}
        self.factor_levels = {}
        self.var_to_cov = {}
        self.var_types = {}
        self.var_mins = {}
        self.var_maxs = {}
        self.subgroup_variables = []
        self.term_names: list[str] = []
        self.linear_terms: list[int] = []
        self.smooth_terms: list[int] = []
        self.ir_smooth_terms: list[int] = []
        self.random_terms: list[int] = []
        self.has_intercept = False
        self.has_irf = False
        self.n_irf = 0
        self.unpenalized_coef: int | None = None
        self.coef_names: list[str] | None = None
        self.n_coef: int | None = None  # Number of total coefficients in formula.
        self.coef_per_term: list[int] | None = (
            None  # Number of coefficients associated with each term
        )
        self.coef_idx_per_term: list[np.ndarray] | None = (
            None  # indices in coef vector per term
        )
        self.built_penalties = False
        self.find_nested = find_nested
        cvi = 0  # Number of variables included in some way as predictors

        # Encoding from data frame to series-level dependent values + predictor values (in cov)
        # sid holds series end indices for quick splitting.
        self.y_flat: np.ndarray | None = None
        self.cov_flat: np.ndarray | None = None
        self.NOT_NA_flat: np.ndarray | None = None
        self.y: list[np.ndarray] | None = None
        self.cov: list[np.ndarray] | None = None
        self.NOT_NA: list[np.ndarray] | None = None
        self.sid: np.ndarray | None = None

        # Discretization?
        self.discretize = {}

        # Perform input checks first for LHS/Dependent variable.
        if len(self.file_paths) == 0 and self.lhs.variable not in self.data.columns:
            raise IndexError(
                f"Column '{self.lhs.variable}' does not exist in Dataframe."
            )

        if len(self.file_paths) != 0 and self.keep_cov is False and self.find_nested:
            if print_warn:
                warnings.warn(
                    (
                        "``find_nested=True`` is not supported when iteratively building the model "
                        "matrix and not keeping the encoded data (i.e., when ``keep_cov=False``). "
                        "Setting ``find_nested=False``"
                    )
                )
            self.find_nested = False

        # Now some checks on the terms - some problems might only be caught later when the
        # penalties are built.
        for ti, term in enumerate(self.terms):

            # Collect term name
            self.term_names.append(term.name)

            # Term allocation.
            if isinstance(term, i):
                self.has_intercept = True
                self.linear_terms.append(ti)
                continue

            if isinstance(term, l):
                self.linear_terms.append(ti)

            if isinstance(term, f):
                self.smooth_terms.append(ti)

            if isinstance(term, irf):
                self.ir_smooth_terms.append(ti)
                self.n_irf += 1

            if isinstance(term, ri) or isinstance(term, rs):
                self.random_terms.append(ti)

            if isinstance(term, fs):
                if term.approx_deriv is not None:
                    self.discretize[ti] = term.approx_deriv

                    # Make sure all categorical split variables end up being encoded since
                    # they do not necessarily have to be in the formula in case of
                    # sub-groups.
                    for split_by_fac in self.discretize[ti]["split_by"]:
                        cvi = self.__encode_var(split_by_fac, "O", cvi, codebook)

            # All variables must exist in data
            for var in term.variables:

                if len(self.file_paths) == 0 and var not in self.data.columns:
                    raise KeyError(
                        f"Variable '{var}' of term {ti} does not exist in dataframe."
                    )

                if len(self.file_paths) == 0:
                    vartype = data[var].dtype
                else:
                    vartype = read_dtype(
                        var, self.file_paths[0], self.file_loading_kwargs
                    )

                # Store information for all variables once.
                cvi = self.__encode_var(var, vartype, cvi, codebook)

                # Smooth-term variables must all be continuous
                if isinstance(term, f) or isinstance(term, irf):
                    if vartype not in ["float64", "int64"]:
                        raise TypeError(
                            (
                                f"Variable '{var}' attributed to smooth/impulse response smooth "
                                f"term {ti} must be numeric and is not."
                            )
                        )

                # Random intercept variable must be categorical
                if isinstance(term, ri):
                    if vartype in ["float64", "int64"]:
                        raise TypeError(
                            (
                                f"Variable '{var}' attributed to random intercept term {ti} must "
                                "not be numeric but is."
                            )
                        )

            # by-variables must be categorical
            if isinstance(term, f) or isinstance(term, irf) or isinstance(term, rs):
                if term.by is not None or (
                    isinstance(term, f) and term.binary is not None
                ):

                    t_by = term.by
                    if t_by is None:
                        t_by = term.binary[0]

                    if len(self.file_paths) == 0 and t_by not in self.data.columns:
                        raise KeyError(
                            (
                                f"By-variable '{t_by}' attributed to term {ti} does not exist in "
                                "dataframe."
                            )
                        )

                    if len(self.file_paths) == 0 and data[t_by].dtype in [
                        "float64",
                        "int64",
                    ]:
                        raise KeyError(
                            (
                                f"Data-type of By-variable '{t_by}' attributed to term {ti} must "
                                "not be numeric but is. E.g., Make sure the pandas dtype is "
                                "'object'."
                            )
                        )

                    if len(self.file_paths) > 0 and read_dtype(
                        t_by, self.file_paths[0], self.file_loading_kwargs
                    ) in ["float64", "int64"]:
                        raise KeyError(
                            (
                                f"Data-type of By-variable '{t_by}' attributed to term {ti} must "
                                "not be numeric but is. E.g., Make sure the pandas dtype is "
                                "'object'."
                            )
                        )

                    t_by_subgroup = None
                    # Handle sub-cluster for factor smooth term.
                    if isinstance(term, fs) and term.by_subgroup is not None:

                        t_by_subgroup = term.by_subgroup

                        if (
                            len(self.file_paths) == 0
                            and not t_by_subgroup[0] in self.data.columns
                        ):
                            raise KeyError(
                                (
                                    f"Sub-group by-variable '{t_by_subgroup}' attributed to term "
                                    f"{ti} does not exist in dataframe."
                                )
                            )

                        if len(self.file_paths) == 0 and data[
                            t_by_subgroup[0]
                        ].dtype in ["float64", "int64"]:
                            raise KeyError(
                                (
                                    f"Data-type of sub-group by-variable '{t_by_subgroup[0]}' "
                                    f"attributed to term {ti} must not be numeric but is. E.g., "
                                    "Make sure the pandas dtype is 'object'."
                                )
                            )

                        if len(self.file_paths) > 0 and read_dtype(
                            t_by_subgroup[0],
                            self.file_paths[0],
                            self.file_loading_kwargs,
                        ) in ["float64", "int64"]:
                            raise KeyError(
                                (
                                    f"Data-type of sub-group by-variable '{t_by_subgroup[0]}' "
                                    f"attributed to term {ti} must not be numeric but is. E.g., "
                                    "Make sure the pandas dtype is 'object'."
                                )
                            )

                        # Make sure sub-group variable is also encoded.
                        cvi = self.__encode_var(t_by_subgroup[0], "O", cvi, codebook)

                    # Store information for by variables as well.
                    cvi = self.__encode_var(
                        t_by, "O", cvi, codebook, by_subgroup=t_by_subgroup
                    )

                    # If no error was raised we can now over-write the by variable
                    # for the factor smooth
                    if isinstance(term, fs) and t_by_subgroup is not None:
                        term.by += ":" + t_by_subgroup[1]

                    if isinstance(term, f) and term.binary is not None:
                        term.binary_level = self.factor_codings[t_by][term.binary[1]]

                # Continuous variable to be multiplied with model matrix for this term.
                if term.by_cont is not None:
                    t_by_cont = term.by_cont

                    if len(self.file_paths) == 0 and t_by_cont not in self.data.columns:
                        raise KeyError(
                            (
                                f"By_cont-variable '{t_by_cont}' attributed to term {ti} does not "
                                "exist in dataframe."
                            )
                        )

                    if len(self.file_paths) == 0:
                        vartype = data[t_by_cont].dtype
                    else:
                        vartype = read_dtype(
                            t_by_cont, self.file_paths[0], self.file_loading_kwargs
                        )

                    if vartype not in ["float64", "int64"]:
                        raise TypeError(
                            (
                                f"Variable '{t_by_cont}' attributed to term {ti} must be numeric "
                                "but is not."
                            )
                        )

                    cvi = self.__encode_var(t_by_cont, vartype, cvi, codebook)

        if self.n_irf > 0:
            self.has_irf = True

        if self.has_irf and (len(self.file_paths) != 0 or self.data is None):
            raise NotImplementedError(
                (
                    "Building X.T@X iteratively does not support Impulse Response Terms "
                    "(i.e., ``irf``) in the formula."
                )
            )

        # Encode data into columns usable by the model
        if len(self.file_paths) == 0 or self.keep_cov:
            y_flat, cov_flat, NAs_flat, y, cov, NAs, sid = self.encode_data(self.data)

            # Store encoding
            self.y_flat = y_flat
            self.cov_flat = cov_flat
            self.NOT_NA_flat = NAs_flat
            self.y = y
            self.cov = cov
            self.NOT_NA = NAs
            self.sid = sid

        if len(self.discretize) > 0:
            if self.series_id is None:
                raise ValueError(
                    (
                        "The identifier column for unique series must be provided when requesting "
                        "to approximate the derivative of one or more factor smooth terms."
                    )
                )

            if len(self.file_paths) != 0 and self.keep_cov is False:
                raise ValueError(
                    (
                        "``Formula.keep.cov`` must be set to ``True`` when reading data from file "
                        "AND approximating the derivative for factor smooths."
                    )
                )

            for sti in self.discretize.keys():
                self.__cluster_discretize(
                    *self.__split_discretize(self.__discretize(sti), sti), sti
                )

            # if len(self.file_paths) != 0 and self.keep_cov == False:
            #   # Clean up
            #   self.cov_flat = None
            #   self.sid = None
            #   self.series_id = None

        # Absorb any constraints for model terms
        if len(self.file_paths) == 0:
            self.__absorb_constraints()
        else:
            self.__absorb_constraints2()

        # Can now check nesting
        if self.find_nested and len(self.smooth_terms) > 0:
            self.__fix_nested()

        # Compute (final) number of coef and coef names
        self.__get_coef_info()

        # print(self.n_coef,len(self.coef_names))

    def __encode_var(
        self,
        var: str,
        vartype: np.dtype,
        cvi: int,
        codebook: dict,
        by_subgroup: tuple[str, str] | None = None,
    ) -> int:
        """Internal function that does bookkeeping on variables, encoding them into the
        ``codebook``.

        :param var: name of variable
        :type var: str
        :param vartype: type of variable
        :type vartype: np.dtype
        :param cvi: variable index
        :type cvi: int
        :param codebook: codebook dictionary
        :type codebook: dict
        :param by_subgroup: Either None or a tuple with two strings: first corresponding to a
            factor, second to a level fo that factor, defaults to None
        :type by_subgroup: tuple[str,str] | None, optional
        :return: Updated cvi
        :rtype: int
        """
        # Store information for all variables once.
        if by_subgroup is not None:
            _org_var = var
            var += ":" + by_subgroup[1]
            self.subgroup_variables.append(var)

        if var not in self.var_to_cov:
            self.var_to_cov[var] = cvi

            # Assign vartype enum and calculate mins/maxs for continuous variables
            if vartype in ["float64", "int64"]:
                # ToDo: these can be properties of the formula.
                self.var_types[var] = VarType.NUMERIC
                if len(self.file_paths) == 0:
                    self.var_mins[var] = np.min(self.data[var])
                    self.var_maxs[var] = np.max(self.data[var])
                else:
                    unique_var = read_unique(
                        var,
                        self.file_paths,
                        self.file_loading_nc,
                        self.file_loading_kwargs,
                    )
                    self.var_mins[var] = np.min(unique_var)
                    self.var_maxs[var] = np.max(unique_var)
            else:
                self.var_types[var] = VarType.FACTOR
                self.var_mins[var] = None
                self.var_maxs[var] = None

                # Code factor variables into integers for easy dummy coding
                if len(self.file_paths) == 0:
                    if by_subgroup is None:
                        levels = np.unique(self.data[var])
                    else:
                        levels = np.unique(
                            self.data.loc[
                                self.data[by_subgroup[0]] == by_subgroup[1], _org_var
                            ]
                        )
                else:
                    if by_subgroup is None:
                        levels = read_unique(
                            var,
                            self.file_paths,
                            self.file_loading_nc,
                            self.file_loading_kwargs,
                        )
                    else:
                        rf_fac = read_cov(
                            self.lhs.variable,
                            _org_var,
                            self.file_paths,
                            self.file_loading_nc,
                            self.file_loading_kwargs,
                        )
                        # print(len(rf_fac))
                        sub_fac = read_cov(
                            self.lhs.variable,
                            by_subgroup[0],
                            self.file_paths,
                            self.file_loading_nc,
                            self.file_loading_kwargs,
                        )
                        # print(len(rf_fac[sub_fac == by_subgroup[1]]))
                        levels = np.unique(rf_fac[sub_fac == by_subgroup[1]])

                self.factor_codings[var] = {}
                self.coding_factors[var] = {}
                self.factor_levels[var] = levels

                for ci, c in enumerate(levels):
                    if codebook is not None and var in codebook:
                        self.factor_codings[var][c] = codebook[var][c]
                        self.coding_factors[var][codebook[var][c]] = c
                    else:
                        self.factor_codings[var][c] = ci
                        self.coding_factors[var][ci] = c

            cvi += 1

        return cvi

    def __get_coef_info(self) -> None:
        """Get's information about the number of coefficients from each term + the names of these
        coefficients. Terms also provide info here about the number of unpenalized coefficients that
        come with the term."""

        var_types = self.var_types
        factor_levels = self.factor_levels
        coding_factors = self.coding_factors

        terms = self.terms
        self.unpenalized_coef = 0
        self.n_coef = 0
        self.coef_names = []
        self.coef_per_term = np.zeros(len(terms), dtype=int)
        self.coef_idx_per_term = [[] for _ in range(len(terms))]

        # Need to enforce dummy coding for linear terms not just when there is an intercept, but
        # also when we have multiple factor variables
        has_constant = self.has_intercept

        for lti in self.get_linear_term_idx():
            # Calculate Coef names for linear terms
            lterm = terms[lti]

            if isinstance(lterm, i):
                # Intercept
                t_total_coef, t_unpenalized_coef, t_coef_names = lterm.get_coef_info()

            else:
                # Linear effects
                t_total_coef, t_unpenalized_coef, t_coef_names = lterm.get_coef_info(
                    has_constant, var_types, factor_levels, coding_factors
                )

                # Check whether this term had a factor variable
                if has_constant is False:
                    for var in lterm.variables:
                        if var_types[var] == VarType.FACTOR:
                            has_constant = True
                            break

            self.coef_names.extend(t_coef_names)
            self.coef_per_term[lti] = t_total_coef
            self.coef_idx_per_term[lti] = np.arange(
                self.n_coef, self.n_coef + t_total_coef
            )
            self.n_coef += t_total_coef
            self.unpenalized_coef += t_unpenalized_coef

        for irsti in self.get_ir_smooth_term_idx():
            # Calculate Coef names for impulse response terms
            irsterm = terms[irsti]

            t_total_coef, t_unpenalized_coef, t_coef_names = irsterm.get_coef_info(
                self.has_intercept, factor_levels
            )

            self.coef_names.extend(t_coef_names)
            self.coef_per_term[irsti] = t_total_coef
            self.coef_idx_per_term[irsti] = np.arange(
                self.n_coef, self.n_coef + t_total_coef
            )
            self.n_coef += t_total_coef
            self.unpenalized_coef += t_unpenalized_coef

        for sti in self.get_smooth_term_idx():
            # Calculate Coef names for smooth terms
            sterm = terms[sti]

            t_total_coef, t_unpenalized_coef, t_coef_names = sterm.get_coef_info(
                factor_levels
            )

            self.coef_names.extend(t_coef_names)
            self.coef_per_term[sti] = t_total_coef
            self.coef_idx_per_term[sti] = np.arange(
                self.n_coef, self.n_coef + t_total_coef
            )
            self.n_coef += t_total_coef
            self.unpenalized_coef += t_unpenalized_coef

        for rti in self.get_random_term_idx():
            # Calculate Coef names for random terms
            rterm = terms[rti]

            if isinstance(rterm, ri):
                t_total_coef, t_unpenalized_coef, t_coef_names = rterm.get_coef_info(
                    factor_levels, coding_factors
                )

            elif isinstance(rterm, rs):
                t_total_coef, t_unpenalized_coef, t_coef_names = rterm.get_coef_info(
                    var_types, factor_levels, coding_factors
                )

            self.coef_names.extend(t_coef_names)
            self.coef_per_term[rti] = t_total_coef
            self.coef_idx_per_term[rti] = np.arange(
                self.n_coef, self.n_coef + t_total_coef
            )
            self.n_coef += t_total_coef
            self.unpenalized_coef += t_unpenalized_coef

    def encode_data(self, data: pd.DataFrame, prediction: bool = False) -> tuple[
        np.ndarray | None,
        np.ndarray,
        np.ndarray | None,
        list[np.ndarray] | None,
        list[np.ndarray] | None,
        list[np.ndarray] | None,
        np.ndarray | None,
    ]:
        """
        Encodes ``data``, which needs to be a ``pd.DataFrame`` and by default
        (if ``prediction==False``) builds an index of which rows in ``data`` are NA in the column
        of the dependent variable described by ``self.lhs``.

        :param data: The data to encode.
        :type data: pd.DataFrame
        :param prediction: Whether or not a NA index and a column for the dependent variable should
            be generated.
        :type prediction: bool, optional
        :return: A tuple with 7 (optional) entries: the dependent variable described by
            ``self.lhs``, the encoded predictor variables as a (N,k) array (number of rows matches
            the number of rows of the first entry returned, the number of columns matches the
            number of k variables present in the formula), an indication for each row whether
            the dependent variable described by ``self.lhs`` is NA, like the first entry but split
            into a list of lists by ``self.series_id``, like the second entry but split into a
            list of lists by ``self.series_id``, ike the third entry but split into a list of lists
            by ``self.series_id``, start and end points for the splits used to split the previous
            three elements (identifying the start and end point of every level of
            ``self.series_id``).
        :rtype: (np.ndarray|None, np.ndarray, np.ndarray|None, list[np.ndarray]|None,
            list[np.ndarray]|None, list[np.ndarray]|None, np.ndarray|None)
        """
        # Build NA index
        if prediction:
            NAs = None
            NAs_flat = None
        else:
            if data is None:
                # read in dep var - without NA correction!
                y_flat = read_cov_no_cor(
                    self.lhs.variable,
                    self.file_paths,
                    self.file_loading_nc,
                    self.file_loading_kwargs,
                )
                NAs_flat = np.isnan(y_flat) == False  # noqa: E712
            else:
                NAs_flat = (
                    np.isnan(data[self.lhs.variable].values) == False  # noqa: E712
                )

        if data is not None:
            if (
                not prediction
                and data.shape[0] != data[NAs_flat].shape[0]
                and self.print_warn
            ):
                missing_p = round(
                    (data.shape[0] - data[NAs_flat].shape[0]) / data.shape[0] * 100,
                    ndigits=2,
                )
                warnings.warn(
                    (
                        f"{data.shape[0] - data[NAs_flat].shape[0]} {self.lhs.variable} values "
                        f"({missing_p}%) are NA."
                    )
                )
            n_y = data.shape[0]
        else:
            n_y = len(y_flat)

        id_col = None
        if self.series_id is not None:
            if data is None:
                id_col = read_cov_no_cor(
                    self.series_id,
                    self.file_paths,
                    self.file_loading_nc,
                    self.file_loading_kwargs,
                )
            else:
                id_col = np.array(data[self.series_id])

        var_map = self.get_var_map()
        n_var = len(var_map)
        var_keys = var_map.keys()
        var_types = self.var_types
        factor_coding = self.factor_codings

        # Collect every series from data frame, make sure to maintain the
        # order of the data frame.
        # Based on: https://stackoverflow.com/questions/12926898
        sid = None
        if self.series_id is not None:
            _, id = np.unique(id_col, return_index=True)
            sid = np.sort(id)

        if prediction:  # For encoding new data
            y_flat = None
            y = None
        else:
            # Collect entire y column
            if data is not None:
                y_flat = np.array(data[self.lhs.variable]).reshape(-1, 1)

            # Then split by seried id
            y = None
            NAs = None
            if self.series_id is not None:
                y = np.split(y_flat, sid[1:])

                # Also split NA index
                NAs = np.split(NAs_flat, sid[1:])

        # Now all predictor variables
        cov_flat = np.zeros((n_y, n_var), dtype=float)

        for c in var_keys:
            if data is None:
                if c in self.subgroup_variables:
                    c_raw = read_cov_no_cor(
                        c.split(":")[0],
                        self.file_paths,
                        self.file_loading_nc,
                        self.file_loading_kwargs,
                    )
                else:
                    c_raw = read_cov_no_cor(
                        c,
                        self.file_paths,
                        self.file_loading_nc,
                        self.file_loading_kwargs,
                    )
            else:
                if c in self.subgroup_variables:
                    c_raw = np.array(data[c.split(":")[0]])
                else:
                    c_raw = np.array(data[c])

            if var_types[c] == VarType.FACTOR:

                c_coding = factor_coding[c]

                # Code factor variable
                if c in self.subgroup_variables:
                    # Set level to -1 which will be ignored later when building the factor smooth.
                    c_code = [c_coding[cr] if cr in c_coding else -1 for cr in c_raw]
                else:
                    c_code = [c_coding[cr] for cr in c_raw]

                cov_flat[:, var_map[c]] = c_code

            else:
                cov_flat[:, var_map[c]] = c_raw

        # Now split cov by series id as well
        cov = None
        if self.series_id is not None:
            cov = np.split(cov_flat, sid[1:], axis=0)

        return y_flat, cov_flat, NAs_flat, y, cov, NAs, sid

    def __discretize(self, sti: int) -> np.ndarray:
        """Internal function to discretize covariates.

        :param sti: Smooth term index pointing to smooth holding a discretization dict
        :type sti: int
        :return: np.ndarray holding discretized covariate values of the data set required by
            ``self.discretize[sti]``.
        :rtype: np.ndarray
        """
        dig_cov_flat = np.zeros_like(self.cov_flat)
        var_types = self.var_types
        var_map = self.get_var_map()

        collected = []

        for var in var_types.keys():
            # Skip variables that should be ignored all together.
            # Useful if one or more continuous variables can be split into
            # categorical factor passed along via "split_by"
            if var in self.discretize[sti]["excl"] or var_types[var] == VarType.FACTOR:
                continue

            if var not in self.discretize[sti]["no_disc"]:
                # Discretize continuous variable into k**0.5 bins, where k is the number of unique
                # values this variable took on in the training set (based on, Wood et al., 2017
                # "Gams for Gigadata").
                values = np.linspace(
                    min(self.cov_flat[:, var_map[var]]),
                    max(self.cov_flat[:, var_map[var]]),
                    int(len(np.unique(self.cov_flat[:, var_map[var]])) ** 0.5),
                )
                dig_cov_flat[:, var_map[var]] = np.digitize(
                    self.cov_flat[:, var_map[var]], values
                )
                collected.append(var_map[var])

            # Also collect continuous variables that should not be discretized
            else:
                dig_cov_flat[:, var_map[var]] = self.cov_flat[:, var_map[var]]
                collected.append(var_map[var])

        return dig_cov_flat[:, collected]

    def __split_discretize(
        self, dig_cov_flat_all: np.ndarray, sti: int
    ) -> tuple[list[np.ndarray], list[np.ndarray]]:
        """Internal function to split dscretized covariate objects.

        :param dig_cov_flat_all: np.ndarray holding discretized covariate values of the data set
            required by ``self.discretize[sti]``.
        :type dig_cov_flat_all: np.ndarray
        :param sti: Smooth term index pointing to smooth holding a discretization dict
        :type sti: int
        :return: tuple holding two lists of np.ndarray. First list has split discretized
            covariate objects per series, second has series id
        :rtype: tuple[list[np.ndarray],list[np.ndarray]]
        """
        var_map = self.get_var_map()
        factor_codings = self.factor_codings

        # Create seried id column in ascending order:
        id_col = np.zeros(dig_cov_flat_all.shape[0], dtype=int)
        id_splits = np.split(id_col, self.sid[1:])
        id_splits = [split + i for i, split in enumerate(id_splits)]
        id_col = np.concatenate(id_splits)

        if not self.terms[sti].by_subgroup is None:
            # Adjust for fact that this factor smooth is fitted for separate level of sub-group.
            sub_group_fact = self.terms[sti].by_subgroup[0]
            sub_group_lvl = self.terms[sti].by_subgroup[1]

            # Build index vector corresponding only to series of sub-group level
            sub_lvl_idx = (
                self.cov_flat[:, var_map[sub_group_fact]]
                == factor_codings[sub_group_fact][sub_group_lvl]
            )

            # Just take what belongs to the sub-group from the disceretized matrix
            dig_cov_flat_all = dig_cov_flat_all[sub_lvl_idx, :]

            # For id col a bit more work is necessary..
            # First take again what belongs to sub-group. Now the problem is that series will
            # no longer go from 0-S. So we have to reset the values in id_col to start from zero and
            # then increment towards the number of series included in this sub-group.
            id_col = id_col[sub_lvl_idx]
            # Split based on indices
            _, id = np.unique(id_col, return_index=True)
            sub_sid = np.sort(id)

            # Reset index
            id_splits = np.split(id_col, sub_sid[1:])
            for idsp, _ in enumerate(id_splits):
                id_splits[idsp][:] = idsp

            # And merge again
            id_col = np.concatenate(id_splits)

        if len(self.discretize[sti]["split_by"]) == 0:
            # Don't actually split
            return [dig_cov_flat_all], [id_col]

        # Now split dig_cov_flat_all per level of combination of all factor variables used for
        # splitting, again correcting for any
        # potential sub-grouping
        if self.terms[sti].by_subgroup is None:
            unq_fact_comb, unq_fact_comb_memb = np.unique(
                self.cov_flat[
                    :, [var_map[fact] for fact in self.discretize[sti]["split_by"]]
                ],
                axis=0,
                return_inverse=True,
            )
        else:
            unq_fact_comb, unq_fact_comb_memb = np.unique(
                self.cov_flat[
                    sub_lvl_idx,
                    [var_map[fact] for fact in self.discretize[sti]["split_by"]],
                ],
                axis=0,
                return_inverse=True,
            )
        # Split series id column per level
        fact_series = [
            id_col[unq_fact_comb_memb == fact] for fact in range(len(unq_fact_comb))
        ]

        # Split dig_cov_flat_all per level
        dig_cov_flats = [
            dig_cov_flat_all[unq_fact_comb_memb == fact, :]
            for fact in range(len(unq_fact_comb))
        ]

        return dig_cov_flats, fact_series

    def __cluster_discretize(
        self, dig_cov_flats: list[np.ndarray], fact_series: list[np.ndarray], sti: int
    ) -> None:
        """Internal function that clusters on the discretized covariate objects.

        Takes input from :func:`__split_discretize`. Sets ``self.discretize[sti]["clust_series"]``
        and ``self.discretize[sti]["clust_weights"]``.

        :param dig_cov_flats: split discretized covariate objects per series
        :type dig_cov_flats: list[np.ndarray]
        :param fact_series: series id array per series
        :type fact_series: list[np.ndarray]
        :param sti: Smooth term index pointing to smooth holding a discretization dict
        :type sti: int
        """
        best_series = None
        best_weights = None
        best_error = None

        iterator = range(self.discretize[sti]["restarts"])
        seed = self.discretize[sti]["seed"]

        if self.print_warn:
            iterator = tqdm(iterator, desc="Clustering", leave=True)

        for rep in iterator:
            clust_max_series = []
            weights = []
            error = 0

            for dig_cov_flat, fact_s in zip(dig_cov_flats, fact_series):

                # Create a simple index vector for each unique series in this factor split.
                _, fact_s_idx = np.unique(fact_s, return_index=True)
                sid_idx = fact_s[fact_s_idx]

                # Now compute the number of unique rows across the discretized matrix.
                dig_cov_flat_unq = np.unique(dig_cov_flat, axis=0)

                # Also compute, for each column, the inverse - telling us for each row in the
                # discretized data to which unique value on the corresponding **variable** it
                # belongs.
                dig_cov_flat_unq_memb = np.zeros_like(dig_cov_flat, dtype=int)

                dig_cov_unq_counts = []

                for vari in range(dig_cov_flat.shape[1]):
                    dig_var_flat_unq, dig_var_flat_unq_memb = np.unique(
                        dig_cov_flat[:, vari], return_inverse=True
                    )

                    if vari > 0:
                        dig_var_flat_unq_memb += dig_cov_unq_counts[-1]

                    dig_cov_unq_counts.append(len(dig_var_flat_unq))
                    dig_cov_flat_unq_memb[:, vari] = dig_var_flat_unq_memb[:]

                # Now we prepare the cluster structure:
                # Every series now gets represented by a row vector with sum(dig_cov_unq_counts)
                # entries
                # Every column in these vectors corresponds to a unique value of an individual
                # discretized (or not) co-variate. The first dig_cov_unq_counts[0] columns
                # correspond to the unique values of covariate 1, the next dig_cov_unq_counts[1]
                # columns correspond to covariate 2, and so on.
                # Each column gets assigned the number of times the corresponding unique value
                # exists for the series to which the vector belongs.
                clust = np.zeros((len(sid_idx), sum(dig_cov_unq_counts)))

                # To compute this we just split the inverse per unique series
                # s_split_unq_memb is a list of 2d arrays, each array corresponding to a series
                # with the number of rows matching the number of observations collected for that
                # series.
                # The number of columns matches the number of collected (and potentially
                # discretized) covariates.
                s_split_unq_memb = np.split(dig_cov_flat_unq_memb, fact_s_idx[1:])

                # Now we flatten that 2d array and collect the unique values and how often they
                # occur. This gives us for every series the indices in the cluster structure that
                # will be non-zero and the value we need to store in each non-zero cell.
                s_split_unq_cnts = [
                    np.unique(s_dig, return_counts=True) for s_dig in s_split_unq_memb
                ]

                # Then loop over each series and add the counts of the corresponding rows to the
                # cluster structure
                for sidx, (udr, cnts) in enumerate(s_split_unq_cnts):
                    clust[sidx, udr] += cnts

                # Use heuristic to determine the number of clusters also used to discretize
                # individual covariates.
                # Then cluster - for estimation this only has to do once before starting the
                # actual fitting routine
                clust_centroids, clust_lab = scp.cluster.vq.kmeans2(
                    clust,
                    int((dig_cov_flat_unq.shape[1] * dig_cov_flat_unq.shape[0]) ** 0.5),
                    minit="++",
                    seed=seed,
                )

                if seed is not None:
                    seed += 1

                # Compute clustering loss, according to scipy docs:
                # https://docs.scipy.org/doc/scipy/reference/generated/scipy.cluster.vq.kmeans2.html
                # Simply pick the cluster set out of all repetitions that minimizes the loss
                for k in range(clust_centroids.shape[0]):
                    error += np.power(
                        clust[clust_lab == k, :] - clust_centroids[k, :], 2
                    ).sum()

                # Find ordering of series ids based on assigned cluster labels
                arg_sort_clust = np.argsort(clust_lab)

                # Sort cluster labels in ascending order for easy split of cluster ordered
                # ids into cluster groups
                sort_clust = clust_lab[arg_sort_clust]
                idx_clust_sort = np.arange(0, len(clust_lab), dtype=int)[arg_sort_clust]

                # Find cluster split points
                _, cid = np.unique(sort_clust, return_index=True)
                csid = np.sort(cid)

                # Now collect all series ids of a particular cluster into a separate array
                idx_grouped = np.split(idx_clust_sort, csid[1:])

                # Find the (ideally, this is done heuristically after all) most complex series in
                # every cluster:
                # Compute the number of unique rows for each series in a cluster, then pick the
                # maximum. Compute complexity weights for all series in the cluster relative to that
                # maximum.
                # These act as a a proxy of how similar each series is to the cluster
                # prototype/maximum.

                for k, clu in enumerate(idx_grouped):
                    clu_sums = np.sum(clust[clu, :], axis=1)
                    clust_max_series.append(sid_idx[clu[np.argmax(clu_sums)]])
                    weights.append(clu_sums / np.max(clu_sums))

                    # clust_max_series.append(sid_idx[clu[np.argmin(clust_distances)]])
                    # clust_rel_distances = clu_sums/clu_sums[np.argmin(clust_distances)]
                    # weights.append(clust_rel_distances/np.max(clust_rel_distances))

            if (rep == 0) or (error < best_error):
                best_series = np.array(clust_max_series)
                best_weights = weights
                best_error = error

        self.discretize[sti]["clust_series"] = best_series
        self.discretize[sti]["clust_weights"] = best_weights

    def __absorb_constraints2(self) -> None:
        """Internal function to absorb identifiability constraints when data is read from file.

        Wood (2017) has an overview over identifiability constraints and different
        reparameterizations for smooth terms.

        References:
           - Wood, S. N. (2017). Generalized Additive Models: An Introduction with R, Second \
            Edition (2nd ed.).

        :raises ValueError: When re-parameterization is requested for identifiable terms.
        :raises NotImplementedError: When a constraint type other than ``QR`` is requested.
        """

        for sti in self.get_smooth_term_idx():

            sterm = self.terms[sti]
            vars = sterm.variables

            if not sterm.is_identifiable:
                if sterm.should_rp > 0:
                    # Reparameterization of marginals was requested - but no identifiability
                    # constraints necessary.
                    for vi in range(len(vars)):

                        if len(vars) > 1:
                            id_nk = sterm.nk[vi]
                        else:
                            id_nk = sterm.nk

                        var_cov_flat = read_cov(
                            self.lhs.variable,
                            vars[vi],
                            self.file_paths,
                            self.file_loading_nc,
                            self.file_loading_kwargs,
                        )

                        matrix_term_v = sterm.basis(
                            var_cov_flat,
                            None,
                            id_nk,
                            min_c=self.var_mins[vars[vi]],
                            max_c=self.var_maxs[vars[vi]],
                            **sterm.basis_kwargs,
                        )

                        sterm.absorb_repara(
                            vi,
                            scp.sparse.csc_array(matrix_term_v),
                            var_cov_flat,
                            min_c=self.var_mins[vars[vi]],
                            max_c=self.var_maxs[vars[vi]],
                        )

                continue

            term_constraint = sterm._const_type
            sterm.Z = []

            if sterm.should_rp > 0:
                raise ValueError(
                    (
                        "Re-parameterizing identifiable terms is currently not supported when "
                        "files are loaded in to build X.T@X incrementally."
                    )
                )

            if term_constraint == ConstType.QR:

                with mp.Pool(processes=10) as pool:
                    C = pool.starmap(
                        _compute_constraint_single_MP,
                        zip(
                            repeat(sterm),
                            repeat(vars),
                            repeat(self.lhs.variable),
                            self.file_paths,
                            repeat(self.var_mins),
                            repeat(self.var_maxs),
                            repeat(self.file_loading_kwargs),
                        ),
                    )

                if sterm.te or len(vars) == 1:
                    C = np.sum(np.array(C), axis=0)
                    Q, _ = scp.linalg.qr(C, pivoting=False, mode="full")
                    sterm.Z.append(Constraint(Q[:, 1:], ConstType.QR))
                else:
                    for vi in range(len(vars)):
                        CVI = np.sum(np.array([cvi[vi] for cvi in C]), axis=0)
                        Q, _ = scp.linalg.qr(CVI, pivoting=False, mode="full")
                        sterm.Z.append(Constraint(Q[:, 1:], ConstType.QR))
            else:
                raise NotImplementedError(
                    (
                        "Only QR constraints are currently supported when files are loaded in to "
                        "build X.T@X incrementally."
                    )
                )

    def __absorb_constraints(self) -> None:
        """Internal function to absorb identifiability constraints and apply selected
        reparameterizations to individual smooth terms.

        Wood (2017) has an overview over identifiability constraints and different
        reparameterizations for smooth terms.

        References:
           - Wood, S. N. (2017). Generalized Additive Models: An Introduction with R, Second \
            Edition (2nd ed.).
        """
        var_map = self.get_var_map()

        for sti in self.get_smooth_term_idx():

            sterm = self.terms[sti]
            vars = sterm.variables

            if not sterm.is_identifiable:
                if sterm.should_rp > 0:
                    # Reparameterization of marginals was requested - but no identifiability
                    # constraints are needed.
                    for vi in range(len(vars)):

                        if len(vars) > 1:
                            id_nk = sterm.nk[vi]
                        else:
                            id_nk = sterm.nk

                        var_cov_flat = self.cov_flat[
                            self.NOT_NA_flat, var_map[vars[vi]]
                        ]

                        matrix_term_v = sterm.basis(
                            var_cov_flat,
                            None,
                            id_nk,
                            min_c=self.var_mins[vars[vi]],
                            max_c=self.var_maxs[vars[vi]],
                            **sterm.basis_kwargs,
                        )

                        sterm.absorb_repara(
                            vi,
                            scp.sparse.csc_array(matrix_term_v),
                            var_cov_flat,
                            min_c=self.var_mins[vars[vi]],
                            max_c=self.var_maxs[vars[vi]],
                        )

                continue

            term_constraint = sterm._const_type
            sterm.Z = []
            matrix_term = None  # for Te basis
            for vi in range(len(vars)):
                # If a smooth term needs to be identifiable I act as if you would have asked
                # for nk+1! so that the identifiable term is of the dimension expected.

                if len(vars) > 1:
                    id_nk = sterm.nk[vi]
                else:
                    id_nk = sterm.nk

                if sterm.te is False:
                    id_nk += 1

                var_cov_flat = self.cov_flat[self.NOT_NA_flat, var_map[vars[vi]]]

                matrix_term_v = sterm.basis(
                    var_cov_flat,
                    None,
                    id_nk,
                    min_c=self.var_mins[vars[vi]],
                    max_c=self.var_maxs[vars[vi]],
                    **sterm.basis_kwargs,
                )

                if sterm.te is False:

                    if term_constraint == ConstType.QR:
                        # Wood (2017) 5.4.1 Identifiability constraints via QR.
                        C = np.sum(matrix_term_v, axis=0).reshape(-1, 1)
                        Q, _ = scp.linalg.qr(C, pivoting=False, mode="full")
                        sterm.Z.append(Constraint(Q[:, 1:], ConstType.QR))
                    elif term_constraint == ConstType.DROP:
                        sterm.Z.append(
                            Constraint(int(matrix_term_v.shape[1] / 2), ConstType.DROP)
                        )
                    elif term_constraint == ConstType.DIFF:
                        sterm.Z.append(
                            Constraint(int(matrix_term_v.shape[1] / 2), ConstType.DIFF)
                        )

                if sterm.should_rp > 0:
                    # Reparameterization of marginals was requested.
                    if sterm.te:
                        XPb = matrix_term_v

                    # For all non-te we have to absorb constraints first
                    elif sterm.Z[vi].type == ConstType.QR:
                        XPb = matrix_term_v @ sterm.Z[vi].Z
                    elif sterm.Z[vi].type == ConstType.DROP:
                        XPb = np.delete(matrix_term_v, sterm.Z[vi].Z, axis=1)
                    elif sterm.Z[vi].type == ConstType.DIFF:
                        # Applies difference re-coding for sum-to-zero coefficients.
                        # Based on smoothCon in mgcv(2017). See constraints.py
                        # for more details.

                        # fmt: off
                        XPb = np.diff(
                            np.concatenate(
                                (
                                    matrix_term_v[
                                        :, sterm.Z[vi].Z : matrix_term_v.shape[1]  # noqa: E203
                                    ],
                                    matrix_term_v[:, : sterm.Z[vi].Z],
                                ),
                                axis=1,
                            )
                        )
                        # fmt: on
                        XPb = np.concatenate(
                            (
                                XPb[:, XPb.shape[1] - sterm.Z[vi].Z :],  # noqa: E203
                                XPb[:, : XPb.shape[1] - sterm.Z[vi].Z],
                            ),
                            axis=1,
                        )

                    sterm.absorb_repara(
                        vi,
                        scp.sparse.csc_array(XPb),
                        self.cov_flat[self.NOT_NA_flat, var_map[vars[vi]]],
                        min_c=self.var_mins[vars[vi]],
                        max_c=self.var_maxs[vars[vi]],
                    )

                if sterm.te:
                    if vi == 0:
                        matrix_term = matrix_term_v
                    else:
                        matrix_term = TP_basis_calc(matrix_term, matrix_term_v)

            # Now deal with te basis
            if sterm.te:
                if term_constraint == ConstType.QR:
                    C = np.sum(matrix_term, axis=0).reshape(-1, 1)
                    Q, _ = scp.linalg.qr(C, pivoting=False, mode="full")
                    sterm.Z.append(Constraint(Q[:, 1:], ConstType.QR))
                elif term_constraint == ConstType.DROP:
                    sterm.Z.append(
                        Constraint(int(matrix_term.shape[1] / 2), ConstType.DROP)
                    )
                elif term_constraint == ConstType.DIFF:
                    sterm.Z.append(
                        Constraint(int(matrix_term.shape[1] / 2), ConstType.DIFF)
                    )

    def __fix_nested(self) -> None:
        """Internal function that identifies nested terms in the Formula, based on section 5.6.3 of
        Wood (2017) and what is implemented in the ``gam.side`` function in ``mgcv``.

        Automatically figures out which coefficients have to be dropped from individual smooth
        terms and assigns the corresponding indices to ``term.drop_coef``.
        These are then automatically removed from the model matrix and penalty built for the
        particular term (i.e., when calling ``term.build_matrix`` and ``term.build_penalty``
        for example via :func:`build_sparse_matrix_from_formula` and :func:`build_penalties`.)

        References:
           - Wood, S. N. (2017). Generalized Additive Models: An Introduction with R, Second \
            Edition (2nd ed.).
           - mgcv gam.side function, see: \
            https://github.com/cran/mgcv/blob/fb7e8e718377513e78ba6c6bf7e60757fc6a32a9/R/mgcv.r#L563
        """
        sti = self.get_smooth_term_idx()
        sti = [ti for ti in sti if isinstance(self.terms[ti], fs) is False]

        term_var = [self.terms[ti].variables for ti in sti]
        n_term_var = [len(tvar) for tvar in term_var]
        max_var = np.max(n_term_var)

        # Check if we have terms with multiple variables
        if max_var == 1:
            return

        # Have smooth terms and potentially some nesting. Now need an extended list of smooth terms,
        # their variable names (taking into account by, binary, and by_cont keywords), and the
        # number of variables to check more carefully for nesting and to then also do something
        # about it.
        term_var = [[] for _ in sti]
        unlist_var = []
        for tii, ti in enumerate(sti):
            sterm = self.terms[ti]
            vars = sterm.variables

            for var in vars:
                ext_vars = []

                if sterm.by is not None:

                    by_levels = self.factor_levels[sterm.by]
                    ext_vars.extend([var + "_" + str(lvl) for lvl in by_levels])
                else:
                    ext_vars.append(var)

                if sterm.binary is not None:
                    ext_vars = [
                        ext_var + "_" + sterm.binary[0] + "_" + sterm.binary[1]
                        for ext_var in ext_vars
                    ]

                if sterm.by_cont is not None:
                    ext_vars = [ext_var + "_" + sterm.by_cont for ext_var in ext_vars]

                term_var[tii].extend(ext_vars)
                unlist_var.extend(ext_vars)

        # Can now safely check for duplicates
        if len(unlist_var) == len(np.unique(unlist_var)):
            return

        # Identify nested depenndencies as discussed in section 5.6.3 of Wood (2017)

        # First build model matrices for terms, making sure to create multiple ones per level of by
        # variable
        # and to account for binary/by_cont keywords.
        Xs = [[] for _ in sti]
        n_term_by = []
        var_map = self.get_var_map()
        cov_flat = self.cov_flat
        for tii, ti in enumerate(sti):
            sterm = self.terms[ti]
            vars = sterm.variables
            for vi in range(len(vars)):

                if len(vars) > 1:
                    id_nk = sterm.nk[vi]
                else:
                    id_nk = sterm.nk

                if sterm.is_identifiable and sterm.te is False:
                    id_nk += 1

                var_cov_flat = cov_flat[self.NOT_NA_flat, var_map[vars[vi]]]

                matrix_term_v = sterm.basis(
                    var_cov_flat,
                    None,
                    id_nk,
                    min_c=self.var_mins[vars[vi]],
                    max_c=self.var_maxs[vars[vi]],
                    **sterm.basis_kwargs,
                )

                if sterm.te is False:
                    # Absorb identifiability constraints into marginals
                    if sterm.is_identifiable:
                        if sterm.Z[vi].type == ConstType.QR:
                            matrix_term_v = matrix_term_v @ sterm.Z[vi].Z
                        elif sterm.Z[vi].type == ConstType.DROP:
                            matrix_term_v = np.delete(
                                matrix_term_v, sterm.Z[vi].Z, axis=1
                            )
                        elif sterm.Z[vi].type == ConstType.DIFF:
                            # Applies difference re-coding for sum-to-zero coefficients.
                            # Based on smoothCon in mgcv(2017). See constraints.py
                            # for more details.

                            # fmt: off
                            matrix_term_v = np.diff(
                                np.concatenate(
                                    (
                                        matrix_term_v[
                                            :, sterm.Z[vi].Z : matrix_term_v.shape[1]  # noqa: E203
                                        ],
                                        matrix_term_v[:, : sterm.Z[vi].Z],
                                    ),
                                    axis=1,
                                )
                            )
                            matrix_term_v = np.concatenate(
                                (
                                    matrix_term_v[
                                        :, matrix_term_v.shape[1] - sterm.Z[vi].Z :  # noqa: E203
                                    ],
                                    matrix_term_v[
                                        :, : matrix_term_v.shape[1] - sterm.Z[vi].Z
                                    ],
                                ),
                                axis=1,
                            )
                            # fmt: on

                # Absorb reparam into marginal
                if sterm.should_rp > 0:
                    matrix_term_v = matrix_term_v @ sterm.RP[vi].C

                # Prepare te/ti
                if vi == 0:
                    matrix_term = matrix_term_v
                else:
                    matrix_term = TP_basis_calc(matrix_term, matrix_term_v)

            # Now can handle te - ti and univariate are ready at this point
            if sterm.te:
                if sterm.is_identifiable:
                    if sterm.Z[0].type == ConstType.QR:
                        matrix_term = matrix_term @ sterm.Z[0].Z

                    elif sterm.Z[0].type == ConstType.DROP:
                        matrix_term = np.delete(matrix_term, sterm.Z[0].Z, axis=1)

                    elif sterm.Z[0].type == ConstType.DIFF:

                        # fmt: off
                        matrix_term = np.diff(
                            np.concatenate(
                                (
                                    matrix_term[:, sterm.Z[0].Z : matrix_term.shape[1]],  # noqa
                                    matrix_term[:, : sterm.Z[0].Z],
                                ),
                                axis=1,
                            )
                        )
                        matrix_term = np.concatenate(
                            (
                                matrix_term[:, matrix_term.shape[1] - sterm.Z[0].Z :],  # noqa: E203
                                matrix_term[:, : matrix_term.shape[1] - sterm.Z[0].Z],
                            ),
                            axis=1,
                        )
                        # fmt: on

                # Currently no further reparam for te implemented based on full matrix

            # Now deal with by, binary, or by_cont
            # Multiply each row of model matrix by value in by_cont
            if sterm.by_cont is not None:
                by_cont_cov = cov_flat[self.NOT_NA_flat, var_map[sterm.by_cont]]
                matrix_term *= by_cont_cov.reshape(-1, 1)

            # Handle optional by keyword
            n_by = 1
            if sterm.by is not None:
                by_levels = self.factor_levels[sterm.by]
                by_cov = cov_flat[self.NOT_NA_flat, var_map[sterm.by]]
                n_by = len(by_levels)

                # Split by cov
                for by_level in range(len(by_levels)):
                    by_cidx = (by_cov == by_level).astype(int)
                    Xs[tii].append(matrix_term * by_cidx.reshape(-1, 1))

            # Handle optional binary keyword
            elif sterm.binary is not None:
                by_cov = cov_flat[self.NOT_NA_flat, var_map[sterm.binary[0]]]
                by_cidx = (by_cov == sterm.binary_level).astype(int)

                Xs[tii].append(matrix_term * by_cidx.reshape(-1, 1))

            else:
                Xs[tii].append(matrix_term)

            n_term_by.append(n_by)

        # At this point we have all the model matrices and can work on actually identifying
        # nested terms.
        osize = 2
        drop = {}
        while osize <= max_var:

            otis = [ti for ti in range(len(sti)) if n_term_var[ti] == osize]

            for otii, oti in enumerate(otis):

                for by_idx in range(n_term_by[oti]):

                    # Get variables of this term
                    otis_var = np.array(term_var[oti])
                    otis_n_var = n_term_var[oti]
                    otis_n_by = n_term_by[oti]

                    if otis_n_by > 1:
                        otis_var = otis_var[
                            np.arange(by_idx, otis_n_var * otis_n_by, otis_n_by)
                        ]

                    # Now first find all terms with less than osize variables...
                    itis = []
                    X1 = []
                    for ti in range(len(sti)):
                        if n_term_var[ti] <= (osize - 1) and np.any(
                            [t_var in otis_var for t_var in term_var[ti]]
                        ):
                            itis.append(ti)

                            X1.append(Xs[ti][by_idx])

                    # and then all previous terms with exactly osize variables
                    if otii > 0 and len(otis) > 1:
                        for ti in otis[:otii]:
                            if np.any([t_var in otis_var for t_var in term_var[ti]]):
                                itis.append(ti)

                                X1.append(Xs[ti][by_idx])

                    # Build combined model matrix
                    if len(X1) == 0:
                        continue

                    X1 = np.concatenate(X1, axis=1)

                    if self.has_intercept:
                        X1 = np.concatenate(
                            [np.ones(X1.shape[0]).reshape(-1, 1), X1], axis=1
                        )

                    # And figure out identifiability issues
                    nid, rank = eigen_solvers.id_dependencies(
                        X1, Xs[oti][by_idx], np.pow(np.finfo(float).eps, 0.5)
                    )
                    if rank < Xs[oti][by_idx].shape[1]:

                        # Delete dropped coef from Xs
                        Xs[oti][by_idx] = np.delete(Xs[oti][by_idx], nid[rank:], axis=1)

                        # and collect them for later model matrix building
                        nid += by_idx * len(nid)

                        if not sti[oti] in drop:
                            drop[sti[oti]] = list(np.sort(nid[rank:]))
                        else:
                            drop[sti[oti]].extend(np.sort(nid[rank:]))

            osize += 1

        for ti in sti:
            if ti in drop:
                self.terms[ti].drop_coef = drop[ti]

    # Getters #

    def get_lhs(self) -> lhs:
        """Get a copy of the ``lhs`` specified for this formula."""
        return copy.deepcopy(self.lhs)

    def get_terms(self) -> list[GammTerm]:
        """Get a copy of the ``terms`` specified for this formula."""
        return copy.deepcopy(self.terms)

    def get_data(self) -> pd.DataFrame:
        """Get a copy of the ``data`` specified for this formula."""
        return copy.deepcopy(self.data)

    def get_factor_codings(self) -> dict:
        """Get a copy of the factor coding dictionary. Keys are factor variables in the data,
        values are dictionaries, where the keys correspond to the levels (str) of the factor and
        the values to their encoded levels (int)."""
        return copy.deepcopy(self.factor_codings)

    def get_coding_factors(self) -> dict:
        """Get a copy of the factor coding dictionary. Keys are factor variables in the data,
        values are dictionaries, where the keys correspond to the encoded levels (int) of the factor
        and the values to their levels (str)."""
        return copy.deepcopy(self.coding_factors)

    def get_var_map(self) -> dict:
        """Get a copy of the var map dictionary. Keys are variables in the data, values their
        column index in the encoded predictor matrix returned by ``self.encode_data``.
        """
        return copy.deepcopy(self.var_to_cov)

    def get_factor_levels(self) -> dict:
        """Get a copy of the factor levels dictionary. Keys are factor variables in the data,
        values are np.arrays holding the unique levels (as str) of the corresponding factor.
        """
        return copy.deepcopy(self.factor_levels)

    def get_var_types(self) -> dict:
        """Get a copy of the var types dictionary. Keys are variables in the data, values are
        either ``VarType.NUMERIC`` for continuous variables or ``VarType.FACTOR`` for categorical
        variables."""
        return copy.deepcopy(self.var_types)

    def get_var_mins(self) -> dict:
        """Get a copy of the var mins dictionary. Keys are variables in the data, values are either
        the minimum value the variable takes on in ``self.data`` for continuous variables or
        ``None`` for categorical variables."""
        return copy.deepcopy(self.var_mins)

    def get_var_maxs(self) -> dict:
        """Get a copy of the var maxs dictionary. Keys are variables in the data, values are either
        the maximum value the variable takes on in ``self.data`` for continuous variables or
        ``None`` for categorical variables."""
        return copy.deepcopy(self.var_maxs)

    def get_var_mins_maxs(self) -> tuple[dict, dict]:
        """Get a tuple containing copies of both the mins and maxs directory. See
        ``self.get_var_mins`` and ``self.get_var_maxs``."""
        return (copy.deepcopy(self.var_mins), copy.deepcopy(self.var_maxs))

    def get_linear_term_idx(self) -> list[int]:
        """Get a copy of the list of indices that identify linear terms in ``self.terms``."""
        return copy.deepcopy(self.linear_terms)

    def get_smooth_term_idx(self) -> list[int]:
        """Get a copy of the list of indices that identify smooth terms in ``self.terms``."""
        return copy.deepcopy(self.smooth_terms)

    def get_ir_smooth_term_idx(self) -> list[int]:
        """Get a copy of the list of indices that identify impulse response terms in
        ``self.terms``."""
        return copy.deepcopy(self.ir_smooth_terms)

    def get_random_term_idx(self) -> list[int]:
        """Get a copy of the list of indices that identify random terms in ``self.terms``."""
        return copy.deepcopy(self.random_terms)

    def get_n_coef(self) -> int:
        """Get the number of coefficients that are implied by the formula."""
        return self.n_coef

    def get_depvar(self) -> np.ndarray:
        """Get a copy of the encoded dependent variable (defined via ``self.lhs``)."""
        return copy.deepcopy(self.y_flat)

    def get_notNA(self) -> np.ndarray:
        """Get a copy of the encoded 'not a NA' vector for the dependent variable (defined via
        ``self.lhs``)."""
        return copy.deepcopy(self.NOT_NA_flat)

    def get_has_intercept(self) -> bool:
        """Does this formula include an intercept or not."""
        return self.has_intercept

    def has_ir_terms(self) -> bool:
        """Does this formula include impulse response terms or not."""
        return self.has_irf

    def get_term_names(self) -> list[str]:
        """Returns a copy of the list with the names of the terms specified for this formula."""
        return copy.deepcopy(self.term_names)

    def get_subgroup_variables(self) -> list:
        """Returns a copy of sub-group variables for factor smooths."""
        return copy.deepcopy(self.subgroup_variables)


def build_penalties(formula) -> list[LambdaTerm]:
    """Function to build all penalty matrices required by a :class:`Formula`.

    The function is called whenever it is needed, but the example below shows you how to use it in
    case you want to extract the penalties directly.

    Examples::

       from mssm.models import *
       from mssmViz.sim import *
       from mssm.src.python.formula import build_penalties

       # Get some data and formula
       Binomdat = sim3(10000,0.1,family=Binomial(),seed=20)
       formula = Formula(lhs("y"),[i(),f(["x0"]),f(["x1"]),f(["x2"]),f(["x3"])],data=Binomdat)

       # Now extract the penalties
       penalties = build_penalties(formula)

       print(penalties)

    :param formula: A Formula
    :type formula: Formula
    :raises KeyError: If an un-penalized irf term is included in the formula after penalized terms.
    :raises KeyError:  If an un-penalized smooth term is included in the formula after penalized
        terms.
    :raises ValueError: If no start index has been defined by the formula. For testing only.
    :return: A list of all penalties (encoded as :class:`LambdaTerm`) required by the formula
    :rtype: list[LambdaTerm]
    """

    col_S = formula.n_coef
    factor_levels = formula.factor_levels
    terms = formula.terms
    penalties = []
    start_idx = formula.unpenalized_coef

    if start_idx is None:
        ValueError("Penalty start index is ill-defined.")

    cur_pen_idx = start_idx
    prev_pen_idx = start_idx

    for irsti in formula.get_ir_smooth_term_idx():

        irsterm = terms[irsti]
        vars = irsterm.variables

        # Calculate nCoef
        n_coef = irsterm.nk

        if len(vars) > 1:
            n_coef = np.prod(irsterm.nk)

        by_levels = None
        if irsterm.by is not None:
            by_levels = factor_levels[irsterm.by]

        if not irsterm.is_penalized:
            if len(penalties) == 0:

                added_not_penalized = n_coef
                if irsterm.by is not None:
                    added_not_penalized *= len(by_levels)
                start_idx += added_not_penalized
                cur_pen_idx = start_idx

                if formula.print_warn:
                    warnings.warn(
                        (
                            f"Impulse response smooth {irsti} is not penalized. Smoothing terms "
                            "should generally be penalized."
                        )
                    )

            else:
                raise KeyError(
                    (
                        f"Impulse response smooth {irsti} is not penalized and placed in the "
                        "formula after penalized terms. Unpenalized terms should be moved to the "
                        "beginning of the formula, ideally behind any linear terms."
                    )
                )

        else:

            for penid in range(len(irsterm.penalty)):

                # Smooth terms can have multiple penalties.
                # In that case the starting index of every subsequent
                # penalty needs to be reset.
                if penid > 0:
                    cur_pen_idx = prev_pen_idx

                penalties, cur_pen_idx = irsterm.build_penalty(
                    irsti, penalties, cur_pen_idx, penid, factor_levels, col_S
                )

        # Keep track of previous penalty starting index
        prev_pen_idx = cur_pen_idx

    for sti in formula.get_smooth_term_idx():

        sterm = terms[sti]
        vars = sterm.variables

        # Calculate nCoef
        if len(vars) > 1:
            n_coef = np.prod(sterm.nk)
        else:
            n_coef = sterm.nk

        by_levels = None
        if sterm.by is not None:
            by_levels = factor_levels[sterm.by]

        if not sterm.is_penalized:
            if len(penalties) == 0:

                added_not_penalized = n_coef
                if sterm.by is not None:
                    added_not_penalized *= len(by_levels)

                start_idx += added_not_penalized

                if formula.print_warn:
                    warnings.warn(
                        (
                            f"Smooth {sti} is not penalized. Smoothing terms should generally be "
                            "penalized."
                        )
                    )

            else:
                raise KeyError(
                    (
                        f"Smooth {sti} is not penalized and placed in the formula after penalized "
                        "terms. Unpenalized terms should be moved to the beginning of the formula, "
                        "ideally behind any linear terms."
                    )
                )

        else:
            S_j_TP_last = None
            TP_last_n = 0
            for penid in range(len(sterm.penalty)):

                # Smooth terms can have multiple penalties.
                # In that case the starting index of every subsequent
                # penalty needs to be reset.
                if penid > 0:
                    cur_pen_idx = prev_pen_idx

                prev_n_pen = len(penalties)
                penalties, cur_pen_idx = sterm.build_penalty(
                    sti, penalties, cur_pen_idx, penid, factor_levels, col_S
                )

                # Add necessary info for derivative approx. for factor smooth penalties
                if isinstance(sterm, fs):
                    if sterm.approx_deriv is not None:
                        penalties[-1].clust_series = formula.discretize[sti][
                            "clust_series"
                        ]
                        penalties[-1].clust_weights = formula.discretize[sti][
                            "clust_weights"
                        ]

                # Optionally include a Null-space penalty - an extra penalty on the
                # function space not regularized by the penalty we just created:
                if sterm.has_null_penalty:

                    is_reparam1 = False
                    n_pen = len(penalties)

                    # Below we distinguish between TP smooths of multiple variables and single
                    # variable smooths. For TP smooths Marra & Wood (2011) suggest to first sum
                    # over the penalties for individual
                    # variables and then computing the null-space for that summed penalty.
                    if len(vars) > 1:

                        # For multivariate smooths, we need to chek for the possibility of
                        # separate penalties for different bylevels. These might not be identical -
                        # in case different coefficients might have been dropped
                        # for each level. Thus we need to keep track of a list of sums; one sum per
                        # by-level and compute the null-space penalty separately per penalty in
                        # those lists.
                        if (
                            (sterm.drop_coef) is not None
                            and (sterm.by is not None)
                            and (sterm.id is None)
                        ):

                            if penid == 0:
                                S_j_TP_last = [
                                    penalties[-(pidx + 1)].S_J.toarray()
                                    for pidx in range(len(by_levels) - 1, -1, -1)
                                ]
                            else:
                                new_TP_last = [
                                    penalties[-(pidx + 1)].S_J.toarray()
                                    for pidx in range(len(by_levels) - 1, -1, -1)
                                ]
                                S_j_TP_last = [
                                    S_j_TP_last[pidx] + new_TP_last[pidx]
                                    for pidx in range(len(new_TP_last))
                                ]

                        # Here all potential penalties belonging to different bylevels are
                        # equivalent. Or we have no bylevels. OR we have an id.
                        # Any of those means that penalties from different variables are equivalent
                        # and we can just maintain a single penalty sum.
                        else:
                            S_j_last = penalties[-1].S_J.toarray()

                            if penid == 0:
                                S_j_TP_last = [S_j_last]
                            else:
                                S_j_TP_last[0] += S_j_last

                        # Keep track of penalties added for tensor smooths, note that
                        # (n_pen - prev_n_pen) will be > 1 for by level smooths
                        TP_last_n += n_pen - prev_n_pen

                        if penid < (
                            len(sterm.penalty) - 1
                        ):  # Not done yet - more penalties to build
                            continue

                        # This holds the final list of penalty sums
                        S_j_last = S_j_TP_last

                        # In the end, update the number of new penalties based on the number of
                        # variables involed for TPs
                        added_pen = int(TP_last_n / len(vars))

                    else:
                        # check for Demmler & Reinsch reparameterization for univariate smooths
                        S_j_last = [penalties[-1].S_J.toarray()]
                        is_reparam1 = (
                            penalties[-1].type == PenType.REPARAM1
                        )  # Only for univariate smooths should this be true.
                        added_pen = n_pen - prev_n_pen

                    # Number of coefficients depends on which type S_j_last is: list with
                    # multiple elements or not?
                    if len(S_j_last) == 1:
                        idk = [S_j_last[0].shape[1]]
                        last_pen_rep = penalties[
                            -1
                        ].rep_sj  # Might have repetition in case of id keyword
                    else:
                        idk = [SJ.shape[1] for SJ in S_j_last]
                        last_pen_rep = (
                            1  # No repetition for independent by-level penalties!
                        )

                    # Now compute the desired Null-space penalties

                    NULL_DIMs = []  # Rank of Nullspace-penalty
                    NULL_idks = []  # Dimension of penalty
                    NULL_S = []  # Nullspace penalty in [dat,row,col]
                    NULL_D = []  # Root of Nullspace penalty in [dat, row, col]

                    if is_reparam1 is False:
                        # Based on: Marra & Wood (2011) and:
                        # https://rdrr.io/cran/mgcv/man/gam.selection.html
                        # and:
                        # https://eric-pedersen.github.io/mgcv-esa-workshop/slides/03-model-selection.pdf # noqa

                        for sji, S_ji_last in enumerate(S_j_last):
                            s, U = scp.linalg.eigh(S_ji_last)
                            DNULL = U[:, s <= 1e-7]
                            NULL_DIM = DNULL.shape[1]  # Null-space dimension
                            DNULL = DNULL.reshape(S_ji_last.shape[1], -1)

                            SNULL = DNULL @ DNULL.T

                            SNULL = scp.sparse.csc_array(SNULL)
                            DNULL = scp.sparse.csc_array(DNULL)

                            # Data in S and D is in canonical format, for competability this is
                            # translated to data, rows, columns
                            pen_data, pen_rows, pen_cols = translate_sparse(SNULL)
                            chol_data, chol_rows, chol_cols = translate_sparse(DNULL)

                            NULL_S.append([pen_data, pen_rows, pen_cols])
                            NULL_D.append([chol_data, chol_rows, chol_cols])
                            NULL_DIMs.append(NULL_DIM)
                            NULL_idks.append(idk[sji])

                    else:
                        # Under the Demmler & Reinsch (1975) re-parameterization the last
                        # S.shape[1] - S.rank cols/rows correspond to functions in the kernel.
                        # Hence we can simply place identity penalties on those to shrink them
                        # to zero. In this form we can also readily
                        # have separate penalties on different null-space functions! This is how
                        # mgcv implements factor smooths as well.

                        # Can only happen to univariate smooth, so indexing S_j_last is safe here.

                        NULL_DIM = (
                            S_j_last[0].shape[1] - sterm.RP[0].rank
                        )  # Null-space dimension

                        for nci in range(
                            S_j_last[0].shape[1] - NULL_DIM, S_j_last[0].shape[1]
                        ):

                            NULL_S.append([[1], [nci], [nci]])
                            NULL_D.append([[1], [nci], [nci]])
                            NULL_DIMs.append(1)
                            NULL_idks.append(idk[0])

                    # Now iterate over Null-space penalties
                    for nri in range(len(NULL_DIMs)):

                        # Get current Nullspace penalty
                        nri_rank = NULL_DIMs[nri]
                        pen_data, pen_rows, pen_cols = NULL_S[nri]
                        chol_data, chol_rows, chol_cols = NULL_D[nri]
                        idk = NULL_idks[nri]

                        cur_pen_idx = prev_pen_idx

                        lTerm = LambdaTerm(
                            start_index=cur_pen_idx, type=PenType.NULL, term=sti
                        )

                        # Embed first penalty
                        lTerm.D_J_emb, _ = embed_in_S_sparse(
                            chol_data,
                            chol_rows,
                            chol_cols,
                            lTerm.D_J_emb,
                            col_S,
                            idk,
                            cur_pen_idx,
                        )
                        lTerm.S_J_emb, cur_pen_idx = embed_in_S_sparse(
                            pen_data,
                            pen_rows,
                            pen_cols,
                            lTerm.S_J_emb,
                            col_S,
                            idk,
                            cur_pen_idx,
                        )
                        lTerm.S_J = embed_in_Sj_sparse(
                            pen_data, pen_rows, pen_cols, lTerm.S_J, idk
                        )
                        lTerm.rank = nri_rank

                        # Update prev_pen_idx for split by penalties and append current one
                        if len(S_j_last) > 1:
                            penalties.append(lTerm)
                            prev_pen_idx = cur_pen_idx

                        # Handle equal penalties per by-level with id keyword
                        else:
                            # Single penalty added - but could involve by keyword
                            if (added_pen) == 1:

                                # Handle any By-keyword
                                if last_pen_rep > 1:
                                    pen_iter = last_pen_rep - 1

                                    chol_rep = np.tile(chol_data, pen_iter)
                                    idx_row_rep = (
                                        np.repeat(np.arange(pen_iter), len(chol_rows))
                                        * idk
                                    )
                                    idx_col_rep = (
                                        np.repeat(np.arange(pen_iter), len(chol_cols))
                                        * idk
                                    )
                                    chol_rep_row = (
                                        np.tile(chol_rows, pen_iter) + idx_row_rep
                                    )
                                    chol_rep_cols = (
                                        np.tile(chol_cols, pen_iter) + idx_col_rep
                                    )

                                    lTerm.D_J_emb, _ = embed_in_S_sparse(
                                        chol_rep,
                                        chol_rep_row,
                                        chol_rep_cols,
                                        lTerm.D_J_emb,
                                        col_S,
                                        idk * pen_iter,
                                        cur_pen_idx,
                                    )

                                    pen_rep = np.tile(pen_data, pen_iter)
                                    idx_row_rep = (
                                        np.repeat(np.arange(pen_iter), len(pen_rows))
                                        * idk
                                    )
                                    idx_col_rep = (
                                        np.repeat(np.arange(pen_iter), len(pen_cols))
                                        * idk
                                    )
                                    pen_rep_row = (
                                        np.tile(pen_rows, pen_iter) + idx_row_rep
                                    )
                                    pren_rep_cols = (
                                        np.tile(pen_cols, pen_iter) + idx_col_rep
                                    )

                                    lTerm.S_J_emb, cur_pen_idx = embed_in_S_sparse(
                                        pen_rep,
                                        pen_rep_row,
                                        pren_rep_cols,
                                        lTerm.S_J_emb,
                                        col_S,
                                        idk * pen_iter,
                                        cur_pen_idx,
                                    )

                                    lTerm.rep_sj = last_pen_rep

                                # In any case, term can be appended here.
                                lTerm.rank = nri_rank * last_pen_rep
                                penalties.append(lTerm)
                            else:
                                # Independent penalties via by
                                # Append penalty for first level
                                penalties.append(lTerm)

                                # And add the penalties again for the remaining
                                # levels as separate terms
                                for _ in range((added_pen) - 1):
                                    lTerm = LambdaTerm(
                                        start_index=cur_pen_idx,
                                        type=PenType.NULL,
                                        term=sti,
                                    )

                                    # Embed penalties
                                    lTerm.D_J_emb, _ = embed_in_S_sparse(
                                        chol_data,
                                        chol_rows,
                                        chol_cols,
                                        lTerm.D_J_emb,
                                        col_S,
                                        idk,
                                        cur_pen_idx,
                                    )
                                    lTerm.S_J_emb, cur_pen_idx = embed_in_S_sparse(
                                        pen_data,
                                        pen_rows,
                                        pen_cols,
                                        lTerm.S_J_emb,
                                        col_S,
                                        idk,
                                        cur_pen_idx,
                                    )
                                    lTerm.S_J = embed_in_Sj_sparse(
                                        pen_data, pen_rows, pen_cols, lTerm.S_J, idk
                                    )
                                    lTerm.rank = nri_rank
                                    penalties.append(lTerm)

        # Keep track of previous penalty starting index
        prev_pen_idx = cur_pen_idx

    for rti in formula.get_random_term_idx():
        # Build penalties for random terms
        rterm = terms[rti]

        penalties, cur_pen_idx = rterm.build_penalty(
            rti, penalties, cur_pen_idx, factor_levels, col_S
        )

    if cur_pen_idx != col_S:
        raise ValueError(
            (
                f"Penalty dimension {cur_pen_idx},{cur_pen_idx} does not match outer model "
                f"matrix dimension {col_S}"
            )
        )

    formula.built_penalties = True
    return penalties


def build_sparse_matrix_from_formula(
    terms: list[GammTerm],
    has_intercept: bool,
    ltx: list[int],
    irstx: list[int],
    stx: list[int],
    rtx: list[int],
    var_types: dict,
    var_map: dict,
    var_mins: dict,
    var_maxs: dict,
    factor_levels: dict,
    cov_flat: np.ndarray,
    cov: np.ndarray | None,
    pool: mp.pool.Pool | None = None,
    use_only: list[int] | None = None,
    tol: float = 0,
) -> scp.sparse.csc_array:
    """Build model matrix from formula properties.

    This function is used internally to construct model matrices from :class:`Formula` objects.
    For greater convenience see the :func:`build_model_matrix` function.

    **Important**, make sure to only ever call this when ``formula.built_penalties==True`` - see
    the :func:`build_model_matrix` function description.

    :param terms: List of terms of a :class:`Formula`
    :type terms: list[GammTerm]
    :param has_intercept: Indicator of whether the Formula has an intercept or not
    :type has_intercept: bool
    :param ltx: Linear term indices
    :type ltx: list[int]
    :param irstx: Impulse response function term indices
    :type irstx: list[int]
    :param stx: Smooth term indices
    :type stx: list[int]
    :param rtx: Random term indices
    :type rtx: list[int]
    :param var_types: Dictionary holding variable types
    :type var_types: dict
    :param var_map: Dictionary mapping variable names to column indices in the encoded data
    :type var_map: dict
    :param var_mins: Dictionary with variable minimums
    :type var_mins: dict
    :param var_maxs: Dictionary with variable maximums
    :type var_maxs: dict
    :param factor_levels: Dictionary with levels associated with each factor
    :type factor_levels: dict
    :param cov_flat: Encoded data
    :type cov_flat: np.ndarray
    :param cov: Encoded data split by levels of the factor in ``Formula.series_id``
    :type cov: np.ndarray | None, optional
    :param pool: An instance of a multiprocessing pool, defaults to None
    :type pool: mp.pool.Pool | None, optional
    :param use_only: A list of indices corresponding to which terms should actually be built. If
        ``None``, then all terms are build. Terms not built are set to zero columns, defaults to
        None
    :type use_only: list[int] | None, optional
    :param tol: Optional tolerance. Absolute values in the model matrix smaller than this are set
        to actual zeroes, defaults to 0
    :type tol: float, optional
    :return: The model matrix implied by a :class:`Formula`  and ``cov_flat``.
    :rtype: scp.sparse.csc_array
    """
    n_y = cov_flat.shape[0]
    elements = []
    rows = []
    cols = []
    ridx = np.array([ri for ri in range(n_y)])  # ToDo: dtype=int?

    ci = 0

    # Need to enforce dummy coding for linear terms not just when there is an intercept, but also
    # when we have multiple factor variables
    has_constant = has_intercept

    for lti in ltx:
        # Build matrix for linear terms
        lterm = terms[lti]

        if isinstance(lterm, i):
            # Intercept
            new_elements, new_rows, new_cols, new_ci = lterm.build_matrix(
                ci, lti, ridx, use_only
            )

        else:
            # Linear term
            new_elements, new_rows, new_cols, new_ci = lterm.build_matrix(
                has_constant,
                ci,
                lti,
                var_map,
                var_types,
                factor_levels,
                ridx,
                cov_flat,
                use_only,
            )

            # Check whether this term had a factor variable
            if has_constant is False:
                for var in lterm.variables:
                    if var_types[var] == VarType.FACTOR:
                        has_constant = True
                        break

        elements.extend(new_elements)
        rows.extend(new_rows)
        cols.extend(new_cols)
        ci += new_ci

    for irsti in irstx:
        # Impulse response terms need to be calculate for every series individually - costly
        irsterm = terms[irsti]

        new_elements, new_rows, new_cols, new_ci = irsterm.build_matrix(
            ci,
            irsti,
            var_map,
            var_mins,
            var_maxs,
            factor_levels,
            ridx,
            cov,
            use_only,
            pool,
            tol,
        )

        elements.extend(new_elements)
        rows.extend(new_rows)
        cols.extend(new_cols)
        ci += new_ci

    for sti in stx:

        sterm = terms[sti]

        new_elements, new_rows, new_cols, new_ci = sterm.build_matrix(
            ci,
            sti,
            var_map,
            var_mins,
            var_maxs,
            factor_levels,
            ridx,
            cov_flat,
            use_only,
            tol,
        )

        elements.extend(new_elements)
        rows.extend(new_rows)
        cols.extend(new_cols)
        ci += new_ci

    for rti in rtx:
        rterm = terms[rti]

        if isinstance(rterm, ri):
            new_elements, new_rows, new_cols, new_ci = rterm.build_matrix(
                ci, rti, var_map, factor_levels, ridx, cov_flat, use_only
            )

        elif isinstance(rterm, rs):
            new_elements, new_rows, new_cols, new_ci = rterm.build_matrix(
                ci, rti, var_map, var_types, factor_levels, ridx, cov_flat, use_only
            )

        elements.extend(new_elements)
        rows.extend(new_rows)
        cols.extend(new_cols)
        ci += new_ci

    mat = scp.sparse.csc_array((elements, (rows, cols)), shape=(n_y, ci))

    return mat


def build_model_matrix(
    formula: Formula,
    pool: mp.pool.Pool | None = None,
    use_only: list[int] | None = None,
    tol: float = 0,
) -> scp.sparse.csc_array:
    """Function to build the model matrix implied by ``formula``.

    **Important:** A small selection of smooth terms, requires that the penalty matrices are built
    at least once before the model matrix can be build.
    For this reason, you generally must call ``build_penalties(formula)`` before calling
    ``build_model_matrix(formula)`` (interally, mssm checks whether
    ``formula.built_penalties==True``.). See the example below.

    Examples::

       from mssm.models import *
       from mssmViz.sim import *
       from mssmViz.plot import *
       import matplotlib.pyplot as plt

       from mssm.src.python.formula import build_penalties,build_model_matrix

       # Get some data and formula
       Binomdat = sim3(10000,0.1,family=Binomial(),seed=20)
       formula = Formula(lhs("y"),[i(),f(["x0"]),f(["x1"]),f(["x2"]),f(["x3"])],data=Binomdat)

       # First extract the penalties
       penalties = build_penalties(formula)

       # Then the model matrix:
       X = build_model_matrix(formula)

    :param formula: A Formula
    :type formula: Formula
    :param pool: An instance of a multiprocessing pool, defaults to None
    :type pool: mp.pool.Pool | None, optional
    :param use_only: A list of indices corresponding to which terms should actually be built. If
        ``None``, then all terms are build. Terms not built are set to zero columns, defaults
        to None
    :type use_only: list[int] | None, optional
    :param tol: Optional tolerance. Absolute values in the model matrix smaller than this are set
        to actual zeroes, defaults to 0
    :type tol: float, optional
    :raises ValueError: If ``formula.built_penalties == False`` - i.e., it is required that
        ``build_penalties(formula)`` was called before calling ``build_model_matrix(formula)``.
    :raises NotImplementedError: If the ``formula`` was set up to read data from file, rather
        than from a pd.Dataframe.
    :return: The model matrix implied by a :class:`Formula`  and ``cov_flat``.
    :rtype: scp.sparse.csc_array
    """

    if formula.built_penalties is False:
        raise ValueError(
            (
                "You must call ``build_penalties(formula)`` once before calling "
                "``build_model_matrix(formula)``."
            )
        )

    if len(formula.file_paths) != 0:
        raise NotImplementedError(
            (
                "Cannot return the model-matrix if data was read directly from file, rather than "
                "provided as a pd.DataFrame."
            )
        )

    # Get all the objects from formula required
    terms = formula.terms
    has_intercept = formula.has_intercept
    ltx = formula.get_linear_term_idx()
    irstx = formula.get_ir_smooth_term_idx()
    stx = formula.get_smooth_term_idx()
    rtx = formula.get_random_term_idx()
    var_types = formula.get_var_types()
    var_map = formula.get_var_map()
    var_mins = formula.get_var_mins()
    var_maxs = formula.get_var_maxs()
    factor_levels = formula.get_factor_levels()

    cov_flat = formula.cov_flat[formula.NOT_NA_flat]

    if len(irstx) > 0:
        cov_flat = formula.cov_flat  # Need to drop NA rows **after** building!
        cov = formula.cov
    else:
        cov = None

    mat = build_sparse_matrix_from_formula(
        terms,
        has_intercept,
        ltx,
        irstx,
        stx,
        rtx,
        var_types,
        var_map,
        var_mins,
        var_maxs,
        factor_levels,
        cov_flat,
        cov,
        pool,
        use_only,
        tol,
    )

    if len(irstx) > 0:
        mat = mat[formula.NOT_NA_flat, :]

    return mat
