"""Least-squares and robust regression helpers."""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray


def fitols(y: NDArray[np.float64], x: NDArray[np.float64]) -> NDArray[np.float64]:
    """Fit an ordinary least-squares regression without an intercept.

    Parameters
    ----------
    y
        Response vector with shape ``(n_samples,)``.
    x
        Design matrix with shape ``(n_samples, n_features)``.

    Returns
    -------
    NDArray[np.float64]
        Estimated coefficient vector with shape ``(n_features,)``.
    """

    coef, _, _, _ = np.linalg.lstsq(x, y, rcond=None)
    return coef.astype(float, copy=False)


def fitrlm(y: NDArray[np.float64], x: NDArray[np.float64]) -> NDArray[np.float64]:
    """Fit a robust Huber M-regression without an intercept.

    Parameters
    ----------
    y
        Response vector with shape ``(n_samples,)``.
    x
        Design matrix with shape ``(n_samples, n_features)``.

    Returns
    -------
    NDArray[np.float64]
        Robustly estimated coefficient vector with shape ``(n_features,)``.
    """

    import statsmodels.api as sm

    model = sm.RLM(y, x, M=sm.robust.norms.HuberT())
    fit = model.fit()
    return np.asarray(fit.params, dtype=float)
