"""
Linear models. This module is in very early development and is subject to frequent breaking changes.

This module requires the NumPy package. PDS only requires Polars, but you can get all the optional dependencies by

`pip install polars_ds[all]`

"""

# Currently skipping tests for this module because the underlying functions are all tested in one way or
# another.

from __future__ import annotations

import polars as pl
import numpy as np
from typing import List, Tuple, Literal
from .typing import LRSolverMethods, NullPolicy, PolarsFrame, TypeAlias

from polars_ds._polars_ds import PyLR, PyGLM, PyElasticNet, PyOnlineLR

import sys

if sys.version_info >= (3, 11):
    from typing import Self
else:  # 3.10, 3.9, 3.8
    from typing_extensions import Self


# --------------------------------------------------------------------------------


def _handle_nulls_in_df(
    df: PolarsFrame, features: List[str], target: str, null_policy: NullPolicy
) -> PolarsFrame:
    if null_policy == "ignore":
        return df
    elif null_policy == "raise":
        total_null_count = (
            df.lazy().select(*features, target).null_count().collect().sum_horizontal()[0]
        )
        if total_null_count > 0:
            raise ValueError("Nulls found in Dataframe.")
        return df
    elif null_policy == "skip":
        return df.drop_nulls(subset=features + [target])
    elif null_policy == "zero":
        return df.with_columns(pl.col(features).fill_null(0.0)).drop_nulls(subset=target)
    elif null_policy == "one":
        return df.with_columns(pl.col(features).fill_null(1.0)).drop_nulls(subset=target)
    else:
        try:
            fill_value = float(null_policy)
            if np.isfinite(fill_value):
                return df.with_columns(pl.col(features).fill_null(fill_value)).drop_nulls(
                    subset=target
                )
            raise ValueError("When null_policy is a number, it cannot be nan or infinite.")
        except Exception as e:
            raise ValueError(f"Unknown null_policy. Error: {e}")


def _handle_nans_in_np(
    X: np.ndarray,  # N x M
    y: np.ndarray,  # N x 1
    null_policy: NullPolicy,
) -> Tuple[np.ndarray, np.ndarray]:
    if null_policy == "ignore":
        return (X, y)
    elif null_policy == "raise":
        if np.any(np.isnan(X)) | np.any(np.isnan(y)):
            raise ValueError("Nulls found in X or y.")
        return (X, y)
    elif null_policy == "skip":
        row_has_nan = np.any(np.isnan(X), axis=1) | np.any(np.isnan(y), axis=1)
        return (X[~row_has_nan], y[~row_has_nan])
    elif null_policy == "zero":
        y_nans = np.any(np.isnan(y), axis=1)
        return (np.nan_to_num(X, nan=0.0)[~y_nans], y[~y_nans])
    elif null_policy == "one":
        y_nans = np.any(np.isnan(y), axis=1)
        return (np.nan_to_num(X, nan=1.0)[~y_nans], y[~y_nans])
    else:
        try:
            fill_value = float(null_policy)
            if np.isfinite(fill_value):
                y_nans = np.any(np.isnan(y), axis=1)
                return (np.nan_to_num(X, nan=fill_value)[~y_nans], y[~y_nans])
            raise ValueError("When null_policy is a number, it cannot be nan or infinite.")
        except Exception as e:
            raise ValueError(f"Unknown null_policy. Error: {e}")


# --------------------------------------------------------------------------------


class LR:
    """
    Normal or Ridge Regression.
    """

    def __init__(
        self,
        has_bias: bool = False,
        lambda_: float = 0.0,
        solver: LRSolverMethods = "qr",
        feature_names_in_: List[str] | None = None,
    ):
        """
        Parameters
        ----------
        lambda_
            The regularization parameters for ridge. If this is positive, then this class will solve Ridge.
        solver
            Use one of 'svd', 'cholesky' and 'qr' method to solve the least square equation. Default is 'qr'.
        has_bias
            Whether to add a bias term. Also known as intercept in other packages.
        feature_names_in_
            Names for the incoming features, if available. If None, the names will be empty. They will be
            learned if .fit_df() is run later, or .set_input_features() is set later.
        """
        self._lr = PyLR(solver, lambda_, has_bias)
        self.feature_names_in_: List[str] = (
            [] if feature_names_in_ is None else list(feature_names_in_)
        )

    @classmethod
    def from_values(
        cls, coeffs: List[float], bias: float = 0.0, feature_names_in_: List[str] | None = None
    ) -> Self:
        """
        Constructs a LR class instance from coefficients and bias values.

        Parameters
        ----------
        coeffs
            Iterable of numbers representing the coefficients
        bias
            Value for the bias term
        feature_names_in_
            Names for the incoming features, if available. If None, the names will be empty. They will be
            learned if .fit_df() is run later, or .set_input_features() is set later.
        """
        coefficients = np.ascontiguousarray(coeffs, dtype=np.float64).flatten()
        lr = cls(
            has_bias=(bias != 0.0),
            lambda_=0.0,
            solver="Not Solved",
            feature_names_in_=feature_names_in_,
        )
        lr._lr.set_coeffs_and_bias(coefficients, bias)
        return lr

    def is_fit(self) -> bool:
        return self._lr.is_fit()

    def __repr__(self) -> str:
        if self._lr.lambda_ > 0.0:
            output = "Linear Regression (Ridge) Model\n"
        else:
            output = "Linear Regression Model\n"

        if self._lr.is_fit():
            output += f"Coefficients: {list(round(x, 5) for x in self._lr.coeffs)}\n"
            output += f"Bias/Intercept: {self._lr.bias}\n"
        else:
            output += "Not fitted yet."
        return output

    def set_input_features(self, features: List[str]) -> Self:
        """
        Sets the names of input features.

        Parameters
        ----------
        features
            List of strings.
        """
        self.feature_names_in_.clear()
        self.feature_names_in_ = list(features)
        return self

    def coeffs(self) -> np.ndarray:
        """
        Returns a copy of the coefficients.
        """
        return self._lr.coeffs

    def bias(self) -> float:
        self._lr.bias

    def fit(self, X: np.ndarray, y: np.ndarray, null_policy: NullPolicy = "ignore") -> Self:
        """
        Fit the linear regression model on NumPy data.

        Parameters
        ----------
        X
            The feature Matrix. NumPy 2D matrix only.
        y
            The target data. NumPy array. Must be reshape-able to (-1, 1).
        null_policy: Literal['raise', 'skip', 'zero', 'one', 'ignore']
            One of options shown here, but you can also pass in any numeric string. E.g you may pass '1.25' to mean
            fill nulls with 1.25. If the string cannot be converted to a float, an error will be thrown. Note: if
            the target column has null, the rows with nulls will always be dropped. Null-fill only applies to non-target
            columns. If target has null, then the row will still be dropped.
        """
        self._lr.fit(*_handle_nans_in_np(X, y.reshape((-1, 1)), null_policy))
        return self

    def fit_df(
        self,
        df: PolarsFrame,
        features: List[str],
        target: str,
        null_policy: NullPolicy = "skip",
        show_report: bool = False,
    ) -> Self:
        """
        Fit the linear regression model on a dataframe. This will overwrite previously set feature names.
        The null policy only handles null values in df, not NaN values. It is the user's responsibility to handle
        NaN values if they exist in their pipeline.

        Parameters
        ----------
        df
            Either an eager or a lazy Polars dataframe.
        features
            List of strings of column names.
        target
            The target column's name.
        null_policy: Literal['raise', 'skip', 'zero', 'one', 'ignore']
            One of options shown here, but you can also pass in any numeric string. E.g you may pass '1.25' to mean
            fill nulls with 1.25. If the string cannot be converted to a float, an error will be thrown. Note: if
            the target column has null, the rows with nulls will always be dropped. Null-fill only applies to non-target
            columns. If target has null, then the row will still be dropped.
        show_report
            Whether to print out a regression report. This will duplicate work and will not work for Ridge
            regression. E.g. Nothing will be printed if lambda_ > 0.
        """
        if show_report and self._lr.lambda_ == 0.0:
            from . import query_lstsq_report

            print(
                df.lazy()
                .select(
                    query_lstsq_report(
                        *features,
                        target=target,
                    ).alias("report")
                )
                .unnest("report")
                .collect()
            )

        df2 = (
            _handle_nulls_in_df(df.lazy(), features, target, null_policy)
            .select(*features, target)
            .collect()
        )
        X = df2.select(features).to_numpy()
        y = df2.select(target).to_numpy()
        self.feature_names_in_.clear()
        self.feature_names_in_ = list(features)
        return self.fit(X, y)

    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Returns the prediction of this linear model.

        Parameters
        ----------
        X
            Data to predict on, as a matrix
        """
        return self._lr.predict(X).reshape((-1, 1))

    def predict_df(self, df: PolarsFrame, name: str = "prediction") -> PolarsFrame:
        """
        Computes the prediction of the linear model and append it as a column in the dataframe. If input
        is lazy, output will be lazy.

        Parameters
        ----------
        df
            Either an eager or a lazy Polars dataframe.
        name
            The name of the prediction column
        """
        if len(self.feature_names_in_) <= 0:
            raise ValueError(
                "The linear model is not fitted on a dataframe, or no feature names have been given."
                "Not enough info to predict on a dataframe. Hint: try .fit_df() or .set_input_features()."
            )

        pred = pl.sum_horizontal(
            beta * pl.col(c) for c, beta in zip(self.feature_names_in_, self._lr.coeffs)
        )
        bias = self._lr.bias
        if bias != 0.0:
            pred = pred + bias

        return df.with_columns(pred.alias(name))


class ElasticNet:
    """
    Elastic Net Regression.
    """

    def __init__(
        self,
        l1_reg: float,
        l2_reg: float,
        has_bias: bool = False,
        tol: float = 1e-5,
        max_iter: int = 2000,
        feature_names_in_: List[str] | None = None,
    ):
        """
        Initializes an ElasticNet regressor. This is equivalent to Sklearn's Elastic Net if you set
        alpha and l1_ratio to be: `alpha = l1_reg + l2_reg`, and `l1_ratio = l1_reg / (l1_reg + l2_reg)`.

        Parameters
        ----------
        l1_reg
            The l1 regularization parameters for the elastic net.
        l2_reg
            The l2 regularization parameters for the elastic net.
        has_bias
            Whether to add a bias term. Also known as intercept in other packages.
        tol
            When updates are smaller than tol, the algorithm will stop.
        max_iter
            The max number of iterations the algorithm will run.
        feature_names_in_
            Names for the incoming features, if available. If None, the names will be empty. They will be
            learned if .fit_df() is run later, or .set_input_features() is set later.
        """
        if l1_reg <= 0.0 and l2_reg <= 0.0:
            raise ValueError("Cannot have both l1_reg and l2_reg <= 0.")

        self._en = PyElasticNet(l1_reg, l2_reg, has_bias, tol, max_iter)
        self.feature_names_in_: List[str] = (
            [] if feature_names_in_ is None else list(feature_names_in_)
        )

    @classmethod
    def from_values(
        cls, coeffs: List[float], bias: float = 0.0, feature_names_in_: List[str] | None = None
    ) -> Self:
        """
        Constructs a LR class instance from coefficients and bias values.

        Parameters
        ----------
        coeffs
            Iterable of numbers representing the coefficients
        bias
            Value for the bias term
        feature_names_in_
            Names for the incoming features, if available. If None, the names will be empty. They will be
            learned if .fit_df() is run later, or .set_input_features() is set later.
        """
        coefficients = np.ascontiguousarray(coeffs, dtype=np.float64).flatten()
        elastic_net = cls(
            float("nan"),
            float("nan"),
            has_bias=(bias != 0.0),
            tol=1e-5,
            max_iter=2000,
            feature_names_in_=feature_names_in_,
        )
        elastic_net._en.set_coeffs_and_bias(coefficients, bias)
        return elastic_net

    def is_fit(self) -> bool:
        return self._en.is_fit()

    def __repr__(self) -> str:
        output = f"Elastic Net Model\nl1, l2 regularizers: {self._en.regularizers}\n"
        if self._en.is_fit():
            output += f"Coefficients: {list(round(x, 5) for x in self._en.coeffs)}\n"
            output += f"Bias/Intercept: {self._en.bias}\n"
        else:
            output += "Not fitted yet."
        return output

    def set_input_features(self, features: List[str]) -> Self:
        """
        Sets the names of input features.

        Parameters
        ----------
        features
            List of strings.
        """
        self.feature_names_in_.clear()
        self.feature_names_in_ = list(features)
        return self

    def coeffs(self) -> np.ndarray:
        """
        Returns a copy of the coefficients.
        """
        return self._en.coeffs

    def has_bias(self) -> bool:
        return self._en.has_bias()

    def bias(self) -> float:
        return self._en.bias

    def fit(self, X: np.ndarray, y: np.ndarray, null_policy: NullPolicy = "ignore") -> Self:
        """
        Fit the Elastic Net model on NumPy data.

        Parameters
        ----------
        X
            The feature Matrix. NumPy 2D matrix only.
        y
            The target data. NumPy array. Must be reshape-able to (-1, 1).
        null_policy: Literal['raise', 'skip', 'zero', 'one', 'ignore']
            One of options shown here, but you can also pass in any numeric string. E.g you may pass '1.25' to mean
            fill nulls with 1.25. If the string cannot be converted to a float, an error will be thrown. Note: if
            the target column has null, the rows with nulls will always be dropped. Null-fill only applies to non-target
            columns. If target has null, then the row will still be dropped.
        """
        self._en.fit(*_handle_nans_in_np(X, y.reshape((-1, 1)), null_policy))
        return self

    def fit_df(
        self,
        df: PolarsFrame,
        features: List[str],
        target: str,
        null_policy: NullPolicy = "skip",
    ) -> Self:
        """
        Fit the Elastic Net model on a dataframe. This will overwrite previously set feature names.
        The null policy only handles null values in df, not NaN values. It is the user's responsibility to handle
        NaN values if they exist in their pipeline.

        Parameters
        ----------
        df
            Either an eager or a lazy Polars dataframe.
        features
            List of strings of column names.
        target
            The target column's name.
        null_policy: Literal['raise', 'skip', 'zero', 'one', 'ignore']
            One of options shown here, but you can also pass in any numeric string. E.g you may pass '1.25' to mean
            fill nulls with 1.25. If the string cannot be converted to a float, an error will be thrown. Note: if
            the target column has null, the rows with nulls will always be dropped. Null-fill only applies to non-target
            columns. If target has null, then the row will still be dropped.
        """
        df2 = (
            _handle_nulls_in_df(df.lazy(), features, target, null_policy)
            .select(*features, target)
            .collect()
        )
        X = df2.select(features).to_numpy()
        y = df2.select(target).to_numpy()
        self.feature_names_in_.clear()
        self.feature_names_in_ = list(features)
        return self.fit(X, y)

    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Returns the prediction of this linear model.

        Parameters
        ----------
        X
            Data to predict on, as a matrix
        """
        return self._en.predict(X).reshape((-1, 1))

    def predict_df(self, df: PolarsFrame, name: str = "prediction") -> PolarsFrame:
        """
        Computes the prediction of the linear model and append it as a column in the dataframe. If input
        is lazy, output will be lazy.

        Parameters
        ----------
        df
            Either an eager or a lazy Polars dataframe.
        name
            The name of the prediction column
        """
        if len(self.feature_names_in_) <= 0:
            raise ValueError(
                "The linear model is not fitted on a dataframe, or no feature names have been given."
                "Not enough info to predict on a dataframe. Hint: try .fit_df() or .set_input_features()."
            )

        pred = pl.sum_horizontal(
            beta * pl.col(c) for c, beta in zip(self.feature_names_in_, self._en.coeffs)
        )
        bias = self._en.bias
        if bias != 0.0:
            pred = pred + bias

        return df.with_columns(pred.alias(name))


class OnlineLR:
    """
    Normal or Ridge Online Regression. This doesn't support dataframe inputs.

    Because of implementation details, it is not recommended to set has_bias = True here
    if runtime speed is crucial.

    Null Behaviors:
    1. During the initial fit, no nulls/NaNs should be present
    2. During online updates, if the record has null/NaN, then it will be ignored. Nothing will be updated.
    """

    def __init__(
        self,
        lambda_: float = 0.0,
        has_bias: bool = False,
    ):
        """
        lambda_
            The L2 regularization factor
        has_bias
            Whether this should fit the bias term
        """
        self._lr = PyOnlineLR(lambda_, has_bias)

    @classmethod
    def from_coeffs_bias_inverse(cls, coeffs: List[float], bias: float, inv: np.ndarray) -> Self:
        """
        Constructs an online linear regression instance from coefficients, inverse. This copies
        data.

        Parameters
        ----------
        coeffs
            Iterable of numbers representing the coefficients
        bias
            The bias term
        inv
            2D NumPy matrix representing the inverse of XtX in a regression problem.
        """
        coefficients = np.ascontiguousarray(coeffs, dtype=np.float64).flatten()
        lr = cls(has_bias=(bias > 0.0), lambda_=0.0)
        lr._lr.set_coeffs_bias_inverse(coefficients, bias, inv)
        return lr

    def is_fit(self) -> bool:
        return self._lr.is_fit()

    def __repr__(self) -> str:
        if self._lr.lambda_ > 0.0:
            output = "Online Linear Regression (Ridge) Model\n"
        else:
            output = "Online Linear Regression Model\n"

        if self._lr.is_fit():
            output += f"Coefficients: {list(round(x, 5) for x in self._lr.coeffs)}\n"
            output += f"Bias/Intercept: {self._lr.bias}\n"
        else:
            output += "Not fitted yet."
        return output

    def coeffs(self) -> np.ndarray:
        """
        Returns a copy of the current coefficients.
        """
        return self._lr.coeffs

    def bias(self) -> float:
        self._lr.bias

    def inv(self) -> np.ndarray:
        """
        Returns a copy of the current inverse matrix (inverse of XtX in a linear regression).
        """
        return self._lr.inv

    def fit(self, X: np.ndarray, y: np.ndarray) -> Self:
        """
        Initial Fit for the online linear regression model on NumPy data.

        Parameters
        ----------
        X
            The feature Matrix. NumPy 2D matrix only.
        y
            The target data. NumPy array. Must be reshape-able to (-1, 1).
        """
        if np.any(np.isnan(X)) | np.any(np.isnan(y)):
            raise ValueError(
                "Online regression currently must fit without null for the initial fit."
            )

        self._lr.fit(X, y)
        return self

    def update(self, X: np.ndarray, y: np.ndarray | float, c: float = 1.0) -> Self:
        """
        Updates the online linear regression model with one row of data. If the row contains np.nan,
        it will be ignored.

        Parameters
        ----------
        X
            Either a a 1d array or a 2d array with 1 row. Must be reshapeable to a matrix with 1 row.
        y
            Either a scalar, or a 1d array with 1 element, or a 2d array of size 1x1.
        c
            The middle term (C) in the woodbury matrix identity. A value of 1.0 means we add
            the impact of the new data, and a value of -1.0 means we remove the impact of the
            data. Any other value will `scale` the impact of the data.
        """
        if not self.is_fit():
            raise ValueError("You cannot update before the initial fit of the matrix.")

        x_2d = X.reshape((1, -1))
        if isinstance(y, float):
            y_2d = np.array([[y]])
        else:
            y_2d = y.reshape((1, 1))

        self._lr.update(x_2d, y_2d, c)
        return self

    def predict(self, X: np.ndarray) -> np.ndarray:
        """
        Returns the prediction of this online linear model.

        Parameters
        ----------
        X
            Data to predict on, as a matrix
        """
        return self._lr.predict(X).reshape((-1, 1))


# ------------------------------------------------------------------------------------

GLMFamily: TypeAlias = Literal["gaussian", "normal", "poisson", "binomial", "logistic", "gamma"]
GLMSolver: TypeAlias = Literal["irls", "lbfgs"]
LinkFunction: TypeAlias = Literal["id", "identity", "log", "logit", "inverse"]


class GLM:
    """
    Generalized Linear Models.

    The GLM generalizes linear regression by allowing the linear model to be related to the response
    variable via a link function and by allowing the magnitude of the variance of each measurement to be a
    function of its predicted value.

    Currently, the variance function will be determined by the link function.
    If a family is given, then the canonical link function is used. Here is a mapping between currently
    implemented families and their link functions:

    gaussian / normal    ==>   id (x -> x)
    poisson              ==>   log (x -> ln(x))
    binomial / logistic  ==>   logit (x -> ln(x/(1-x)))
    gamma                ==>   inverse (x -> 1/x)

    Reference
    ---------
    https://en.wikipedia.org/wiki/Generalized_linear_model
    """

    def __init__(
        self,
        add_bias: bool = False,
        solver: GLMSolver = "irls",
        family: GLMFamily = "normal",
        max_iter: int = 100,
        tol: float = 1e-8,
        feature_names_in_: List[str] | None = None,
    ):
        # lambda_: float = 0.0,
        """
        Parameters
        ----------
        family
            One of "gaussian", "normal", "poisson", "binomial", "logistic", "gamma". Note "gaussian" and
            "normal" represent the same family.
        add_bias
            Whether to add a bias term. Also known as intercept in other packages.
        max_iter
            Max number of iterations for the algorithm
        tol
            The tolerance for convergence
        feature_names_in_
            Names for the incoming features, if available. If None, the names will be empty. They will be
            learned if .fit_df() is run later, or .set_input_features() is set later.
        """
        if solver not in ["irls"]:
            raise NotImplementedError

        if max_iter < 1:
            raise ValueError("`max_iter` must be > 1.")

        if family not in ["gaussian", "normal", "poisson", "binomial", "logistic", "gamma"]:
            raise NotImplementedError

        self._glm = PyGLM(
            add_bias=add_bias, family=family, solver=solver, max_iter=max_iter, tol=abs(tol)
        )
        self.feature_names_in_: List[str] = (
            [] if feature_names_in_ is None else list(feature_names_in_)
        )

    def __repr__(self) -> str:
        """
        Shows a textual representation of the GLM.
        """
        return self._glm.describe()

    # @classmethod
    # def from_values(
    #     cls,
    #     coeffs: List[float],
    #     link: LinkFunction,
    #     bias: float = 0.0,
    #     feature_names_in_: List[str] | None = None
    # ) -> Self:
    #     """
    #     Constructs a LR class instance from coefficients and bias values.

    #     Parameters
    #     ----------
    #     coeffs
    #         Iterable of numbers representing the coefficients
    #     link
    #         One of ["id", "log", "logit", "inverse"].
    #     bias
    #         Value for the bias term
    #     feature_names_in_
    #         Names for the incoming features, if available. If None, the names will be empty. They will be
    #         learned if .fit_df() is run later, or .set_input_features() is set later.
    #     """
    #     coefficients = np.ascontiguousarray(coeffs, dtype=np.float64).flatten()
    #     lr = cls(
    #         add_bias=(bias != 0.0),
    #         solver="irls",
    #         feature_names_in_=feature_names_in_,
    #     )
    #     lr._lr.set_coeffs_and_bias(coefficients, bias)
    #     return lr

    def is_fit(self) -> bool:
        return self._glm.is_fit()

    # def __repr__(self) -> str:
    #     if self._lr.lambda_ > 0.0:
    #         output = "Linear Regression (Ridge) Model\n"
    #     else:
    #         output = "Linear Regression Model\n"

    #     if self._lr.is_fit():
    #         output += f"Coefficients: {list(round(x, 5) for x in self._lr.coeffs)}\n"
    #         output += f"Bias/Intercept: {self._lr.bias}\n"
    #     else:
    #         output += "Not fitted yet."
    #     return output

    def set_input_features(self, features: List[str]) -> Self:
        """
        Sets the names of input features.

        Parameters
        ----------
        features
            List of strings.
        """
        self.feature_names_in_.clear()
        self.feature_names_in_ = list(features)
        return self

    def coeffs(self) -> np.ndarray:
        """
        Returns a copy of the coefficients.
        """
        return self._glm.coeffs

    def bias(self) -> float:
        return self._glm.bias

    def fit(self, X: np.ndarray, y: np.ndarray, null_policy: NullPolicy = "ignore") -> Self:
        """
        Fit the GLM model on NumPy data.

        Parameters
        ----------
        X
            The feature Matrix. NumPy 2D matrix only.
        y
            The target data. NumPy array. Must be reshape-able to (-1, 1).
        null_policy: Literal['raise', 'skip', 'zero', 'one', 'ignore']
            One of options shown here, but you can also pass in any numeric string. E.g you may pass '1.25' to mean
            fill nulls with 1.25. If the string cannot be converted to a float, an error will be thrown. Note: if
            the target column has null, the rows with nulls will always be dropped. Null-fill only applies to non-target
            columns. If target has null, then the row will still be dropped.
        """
        self._glm.fit(*_handle_nans_in_np(X, y.astype(np.float64).reshape((-1, 1)), null_policy))
        return self

    def fit_df(
        self,
        df: PolarsFrame,
        features: List[str],
        target: str,
        null_policy: NullPolicy = "skip",
        show_report: bool = False,
    ) -> Self:
        """
        Fit the GLM model on a dataframe. This will overwrite previously set feature names.
        The null policy only handles null values in df, not NaN values. It is the user's responsibility to handle
        NaN values if they exist in their pipeline.

        Parameters
        ----------
        df
            Either an eager or a lazy Polars dataframe.
        features
            List of strings of column names.
        target
            The target column's name.
        null_policy: Literal['raise', 'skip', 'zero', 'one', 'ignore']
            One of options shown here, but you can also pass in any numeric string. E.g you may pass '1.25' to mean
            fill nulls with 1.25. If the string cannot be converted to a float, an error will be thrown. Note: if
            the target column has null, the rows with nulls will always be dropped. Null-fill only applies to non-target
            columns. If target has null, then the row will still be dropped.

        """
        df2 = (
            _handle_nulls_in_df(df.lazy(), features, target, null_policy)
            .select(*features, target)
            .collect()
        )
        X = df2.select(features).to_numpy()
        y = df2.select(target).to_numpy()
        self.feature_names_in_.clear()
        self.feature_names_in_ = list(features)
        return self.fit(X, y)

    def predict(self, X: np.ndarray, linear: bool = False) -> np.ndarray:
        """
        Returns the prediction of this linear model.

        Parameters
        ----------
        X
            Data to predict on, as a matrix
        linear
            If true, return the linear predictor eta instead of the expected value of
            the response variable, E[Y|X].
        """
        return self._glm.predict(X, linear).reshape((-1, 1))

    # def predict_df(self, df: PolarsFrame, name: str = "prediction") -> PolarsFrame:
    #     """
    #     Computes the prediction of the linear model and append it as a column in the dataframe. If input
    #     is lazy, output will be lazy.

    #     Parameters
    #     ----------
    #     df
    #         Either an eager or a lazy Polars dataframe.
    #     name
    #         The name of the prediction column
    #     """
    #     if len(self.feature_names_in_) <= 0:
    #         raise ValueError(
    #             "The linear model is not fitted on a dataframe, or no feature names have been given."
    #             "Not enough info to predict on a dataframe. Hint: try .fit_df() or .set_input_features()."
    #         )

    #     pred = pl.sum_horizontal(
    #         beta * pl.col(c) for c, beta in zip(self.feature_names_in_, self._lr.coeffs)
    #     )
    #     bias = self._lr.bias
    #     if bias != 0.0:
    #         pred = pred + bias

    #     return df.with_columns(pred.alias(name))
