import numpy as np
import pytest
from pamle.likelihoods import Exponential, Logistic, Poisson, gather_blocks, STATUS_OK, STATUS_DIVERGED


def fd_check(model, theta, db, eps=1e-6):
    p = theta.shape[1]
    g, H = model.score_hessian(theta, db)
    gfd = np.array([(model.loglik(theta + eps * np.eye(p)[i], db) - model.loglik(theta - eps * np.eye(p)[i], db))[0] / (2 * eps) for i in range(p)])
    Hfd = np.array([(model.score_hessian(theta + eps * np.eye(p)[i], db)[0] - model.score_hessian(theta - eps * np.eye(p)[i], db)[0])[0] / (2 * eps) for i in range(p)])
    S = np.linalg.inv(-H)
    T = model.third_contraction(theta, db, S)
    Tfd = np.array([np.sum((model.score_hessian(theta + eps * np.eye(p)[i], db)[1] - model.score_hessian(theta - eps * np.eye(p)[i], db)[1])[0] / (2 * eps) * S[0]) for i in range(p)])
    return np.max(np.abs(g[0] - gfd)), np.max(np.abs(H[0] - Hfd)), np.max(np.abs(T[0] - Tfd))


@pytest.mark.parametrize("model,theta0", [(Logistic(4), [0.3, 0.5, -0.7, 0.2]), (Poisson(3), [0.4, 0.3, -0.2]), (Exponential(), [1.7])])
def test_derivatives_against_finite_differences(model, theta0):
    rng = np.random.default_rng(11)
    d = model.simulate(rng, 4000, theta0)
    db = {k: v[None] for k, v in d.items()}
    th = np.asarray(theta0, float)[None] * 0.8 + 0.05
    eg, eh, et = fd_check(model, th, db)
    assert eg < 1e-8 and eh < 1e-7 and et < 1e-6


def test_fit_matches_closed_form_exponential():
    rng = np.random.default_rng(3); ex = Exponential()
    d = ex.simulate(rng, 500, 2.0)
    th, info = ex.fit(d)
    assert abs(th[0] - 1 / d['x'].mean()) < 1e-12 and info['status'] == STATUS_OK


def test_batched_fit_equals_individual_fits():
    rng = np.random.default_rng(4); m = Logistic(3)
    d = m.simulate(rng, 2048, [0.2, 0.8, -0.5])
    db = gather_blocks(d, np.arange(2048).reshape(8, 256))
    th, info = m.fit_batched(db, tol=1e-11)
    assert np.all(info['status'] == STATUS_OK)
    for k in range(8):
        thk, _ = m.fit({kk: v[k] for kk, v in db.items()}, tol=1e-11)
        assert np.allclose(th[k], thk, atol=1e-8)


def test_score_zero_at_mle_tight():
    rng = np.random.default_rng(5); m = Poisson(3)
    d = m.simulate(rng, 5000, [0.5, 0.3, -0.2])
    th, info = m.fit(d, tol=1e-12)
    assert info['grad_max'] < 1e-12


def test_separation_is_flagged():
    """Perfectly separated logistic sample: the MLE does not exist; status must not be OK."""
    m = Logistic(2)
    X = np.hstack([np.ones((40, 1)), np.linspace(-1, 1, 40)[:, None]])
    y = (X[:, 1] > 0).astype(float)
    th, info = m.fit({'X': X, 'y': y}, maxit=200)
    assert info['status'] != STATUS_OK


def test_random_blocks_cover_everything():
    from pamle.partitioning import random_blocks, equal_blocks, skewed_sizes, unequal_blocks
    rng = np.random.default_rng(0)
    b = random_blocks(1000, 10, rng); assert sorted(b.ravel().tolist()) == list(range(1000))
    assert equal_blocks(12, 3).shape == (3, 4)
    s = skewed_sizes(1000, 7, 4.0); assert s.sum() == 1000 and s.max() / s.min() > 3
    ub = unequal_blocks(s); assert sum(len(u) for u in ub) == 1000