import numpy as np
from numpy.testing import assert_allclose
import pytest
from sklearn.datasets import make_low_rank_matrix
from sklearn.linear_model import PoissonRegressor
from sklearn.linear_model._linear_loss import LinearModelLoss


alpha = 0  # unpenalized, even tiny values like 1e-10 does not help.
fit_intercept = True  # if set to False, all tests pass.
model = PoissonRegressor(
    alpha = alpha,
    fit_intercept=fit_intercept,
    tol=1e-12,
    max_iter=1000,
)

# same as glm_dataset in test_glm.py for "wide"
n_samples, n_features = 4, 12
k = n_samples
rng = np.random.RandomState(42)
X = make_low_rank_matrix(
    n_samples=n_samples,
    n_features=n_features,
    effective_rank=k,
    tail_strength=0.1,
    random_state=rng,
)
X[:, -1] = 1  # last columns acts as intercept
U, s, Vt = np.linalg.svd(X, full_matrices=False)

raw_prediction = rng.uniform(low=-3, high=3, size=n_samples)
# minimum norm solution min ||w||_2 such that raw_prediction = X w:
# w = X'(XX')^-1 raw_prediction = V s^-1 U' raw_prediction
coef_unpenalized = Vt.T @ np.diag(1 / s) @ U.T @ raw_prediction

linear_loss = LinearModelLoss(base_loss=model._get_loss(), fit_intercept=True)
sw = np.full(shape=n_samples, fill_value=1 / n_samples)
y = linear_loss.base_loss.link.inverse(raw_prediction)

if fit_intercept:
    X = X[:, :-1]  # remove intercept
    intercept = coef_unpenalized[-1]
    coef = coef_unpenalized[:-1]
else:
    intercept = 0
    coef = coef_unpenalized

model.fit(X, y)

# This is true. So we have a solution.
assert_allclose(model.predict(X), y, rtol=1e-6)

# But it is NOT the minimum norm solution
norm_solution = np.linalg.norm(np.r_[intercept, coef])
norm_model = np.linalg.norm(np.r_[model.intercept_, model.coef_])
print(f"norm_solution = {norm_solution}")
print(f"norm_model    = {norm_model}")
assert norm_model == pytest.approx(norm_solution, rel=1e-4)