import numpy as np

try:
    import jax.numpy as jnp
    from jax import device_put
    from jax.numpy.linalg import inv as jinv
except ImportError:
    pass


def get_beta(X, y, backend="cpu"):
    if backend == "cpu":
        return np.linalg.solve(X.T @ X, X.T @ y)
    device_put(X)
    device_put(y)
    return jnp.linalg.solve(X.T @ X, X.T @ y)


def get_beta_ridge(X, y, lam, backend="cpu"):
    if backend == "cpu":
        return np.linalg.inv(X.T @ X + lam * np.eye(X.shape[1])) @ X.T @ y
    else:
        device_put(X)
        device_put(y)
        XTX = X.T @ X
        Xty = X.T @ y
        I = np.eye(X.shape[1])
        device_put(XTX)
        device_put(Xty)
        device_put(I)
        return jinv(XTX + lam * I) @ Xty


def get_best_beta(X, y, lambdas, backend="cpu"):

    # Number of features
    n_features = X.shape[1]

    # Compute X^T X and X^T y
    if backend == "gpu":
        device_put(X)
        device_put(y)

    Xt = X.T
    XtX = Xt @ X
    Xty = Xt @ y
    I = np.eye(n_features)

    if backend == "gpu":
        device_put(Xt)
        device_put(XtX)
        device_put(Xty)
        device_put(I)
        hat_matrices_diags = [
            (X @ jinv(XtX + lam * I) @ Xt).diagonal() for lam in lambdas
        ]
    else:
        hat_matrices_diags = [
            (X @ np.linalg.inv(XtX + lam * I) @ Xt).diagonal()
            for lam in lambdas
        ]
    train_errors = [
        np.linalg.norm(
            (y - X @ get_beta_ridge(X, y, lam=lam, backend=backend))
            / (1 - hat_matrices_diags[i])
        )
        for i, lam in enumerate(lambdas)
    ]
    train_errors = np.asarray(train_errors)
    best_lam = lambdas[np.argmin(train_errors)]
    # best coefficients
    return (
        get_beta_ridge(X, y, lam=best_lam, backend=backend),
        train_errors,
        best_lam,
    )
