"""
Model Equivalence

Full-stack test: LBFGSStreamingCallback is invoked with the fixture trajectory
(just like scipy.minimize would), computing alpha, s, z internally. ELBO from
the callback must be statistically equivalent to the reference.

Fixture files live in tests/pathfinder/fixtures/ and are generated by running:

    python tests/pathfinder/generate_fixtures.py

Each .npz file contains:
    x_full     – position history  (L+1, N)
    g_full     – gradient history  (L+1, N)
    elbo_ref   – reference ELBO    (L,)

Test strategy
-------------
Statistical equivalence (platform-independent):
    - Spearman rank correlation > 0.75
    - argmax within max(5, L//5) steps of the reference argmax
    - Ballpark: 90% of ELBO values within 20% (relative) of reference
"""

import os

from unittest.mock import MagicMock, patch

import numpy as np
import pymc as pm
import pytest

from pymc.blocking import DictToArrayBijection
from pymc.initial_point import make_initial_point_fn
from pymc.model.core import Point
from pytensor.compile.mode import Mode
from pytensor.tensor.optimize import LRUCache1

from pymc_extras.inference.pathfinder.lbfgs import LBFGSStatus
from pymc_extras.inference.pathfinder.pathfinder import (
    DEFAULT_LINKER,
    FAILED_PATH_STATUS,
    LBFGSStreamingCallback,
    get_neg_logp_dlogp_of_ravel_inputs,
    make_pathfinder_sample_fn,
    make_single_pathfinder_fn,
)
from tests.pathfinder.equivalence_models import MODEL_FACTORIES, make_ard_regression

# Must stay in sync with generate_fixtures.py
FIXTURES_DIR = os.path.join(os.path.dirname(__file__), "fixtures")
ELBO_SEED = 12345
MAXCOR = 6
MAXITER = 100
NUM_DRAWS = 50
NUM_ELBO_DRAWS = 100
JITTER = 2.0
EPSILON = 1e-12

COMPILE_KWARGS = {"mode": Mode(linker=DEFAULT_LINKER)}
ALL_MODEL_NAMES = list(MODEL_FACTORIES.keys())


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _load_fixture(name: str):
    path = os.path.join(FIXTURES_DIR, f"{name}.npz")
    if not os.path.exists(path):
        pytest.skip(
            f"Fixture not found: {path}. Run tests/pathfinder/generate_fixtures.py to create it."
        )
    with np.load(path) as data:
        required = ("x_full", "g_full", "elbo_ref")
        if not all(k in data for k in required):
            pytest.skip(
                f"Fixture {path} is outdated. "
                "Re-run tests/pathfinder/generate_fixtures.py to regenerate."
            )
        return data["x_full"], data["g_full"], data["elbo_ref"]


def _run_callback_replay(
    model: pm.Model,
    x_full: np.ndarray,
    seed: int = ELBO_SEED,
    vectorize: bool = False,
) -> np.ndarray:
    """Replay x_full through LBFGSStreamingCallback (full stack: alpha, s, z computed internally)."""
    neg_logp_dlogp = get_neg_logp_dlogp_of_ravel_inputs(model, jacobian=True, **COMPILE_KWARGS)

    N = x_full.shape[1]
    sample_logp_fn = make_pathfinder_sample_fn(
        model, N, MAXCOR, jacobian=True, compile_kwargs=COMPILE_KWARGS, vectorize=vectorize
    )
    cached_fn = LRUCache1(neg_logp_dlogp, copy_x=True)
    rng = np.random.default_rng(seed)

    elbo_history = []

    def on_step(x, g, alpha, s_win, z_win, elbo):
        elbo_history.append(elbo if np.isfinite(elbo) else np.nan)

    cb = LBFGSStreamingCallback(
        value_grad_fn=cached_fn,
        x0=x_full[0],
        sample_logp_fn=sample_logp_fn,
        num_elbo_draws=NUM_ELBO_DRAWS,
        rng=rng,
        J=MAXCOR,
        epsilon=EPSILON,
        on_step_callback=on_step,
    )

    # Replay trajectory (same as scipy.minimize callback invocation)
    _, g_init = cached_fn(x_full[0])
    cb.x_prev = x_full[0].copy()
    cb.g_prev = np.array(g_init, dtype=np.float64)

    for x_k in x_full[1:]:
        cached_fn(x_k)  # warm cache
        cb(x_k)

    return np.array(elbo_history)


def _make_single_fn(model, maxiter: int = MAXITER, maxcor: int = MAXCOR, vectorize: bool = False):
    return make_single_pathfinder_fn(
        model=model,
        num_draws=NUM_DRAWS,
        maxcor=maxcor,
        maxiter=maxiter,
        ftol=1e-5,
        gtol=1e-8,
        maxls=1000,
        num_elbo_draws=NUM_ELBO_DRAWS,
        jitter=JITTER,
        epsilon=1e-8,
        compile_kwargs=COMPILE_KWARGS,
        pathfinder_kwargs={"vectorize": vectorize},
    )


def _make_lbfgs_mock(x_full):
    """Patch LBFGS to replay x_full[1:] through the streaming callback."""
    mock_inst = MagicMock()

    def replay_streaming(callback, x0):
        _, g_init = callback.value_grad_fn(x_full[0])
        callback.x_prev = x_full[0].copy()
        callback.g_prev = np.array(g_init, dtype=np.float64)

        for x_k in x_full[1:]:
            callback.value_grad_fn(x_k)
            callback(x_k)

        return x_full.shape[0] - 1, LBFGSStatus.CONVERGED

    mock_inst.minimize_streaming.side_effect = replay_streaming

    return patch("pymc_extras.inference.pathfinder.pathfinder.LBFGS", return_value=mock_inst)


def _run_streaming_lbfgs(
    model: pm.Model,
    sample_logp_fn,
    *,
    seed: int,
):
    """Run one streaming L-BFGS optimization and return callback state."""
    neg_logp_dlogp = get_neg_logp_dlogp_of_ravel_inputs(model, jacobian=True, **COMPILE_KWARGS)
    N = DictToArrayBijection.map(model.initial_point()).data.shape[0]

    ipfn = make_initial_point_fn(model=model)
    ip = Point(ipfn(None), model=model)
    x_base = DictToArrayBijection.map(ip).data
    rng = np.random.default_rng(seed)
    x0 = x_base + rng.uniform(-JITTER, JITTER, size=x_base.shape)

    cached_fn = LRUCache1(neg_logp_dlogp, copy_x=True)
    cb = LBFGSStreamingCallback(
        value_grad_fn=cached_fn,
        x0=x0,
        sample_logp_fn=sample_logp_fn,
        num_elbo_draws=NUM_ELBO_DRAWS,
        rng=rng,
        J=MAXCOR,
        epsilon=1e-8,
    )

    from scipy.optimize import minimize as scipy_minimize

    scipy_minimize(
        cached_fn,
        x0,
        method="L-BFGS-B",
        jac=True,
        callback=cb,
        options={"maxcor": MAXCOR, "maxiter": MAXITER, "ftol": 1e-5, "gtol": 1e-8, "maxls": 1000},
    )

    return cb, rng, N


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("vectorize", [False, True])
def test_elbo_call_count(vectorize):
    """sample_logp_fn is called exactly once per accepted LBFGS step."""
    model = make_ard_regression()
    N = DictToArrayBijection.map(model.initial_point()).data.shape[0]

    sample_logp_fn = make_pathfinder_sample_fn(
        model, N, MAXCOR, jacobian=True, compile_kwargs=COMPILE_KWARGS, vectorize=vectorize
    )

    call_count = [0]

    def counting_fn(x, g, alpha, s_win, z_win, u):
        call_count[0] += 1
        return sample_logp_fn(x, g, alpha, s_win, z_win, u)

    cb, _, _ = _run_streaming_lbfgs(model, counting_fn, seed=0)

    assert call_count[0] == cb.step_count
    assert cb.step_count > 0


def test_inv_hessian_diag_matches_sample_scales():
    """Diagonal estimate from state matches empirical sample variances."""
    model = make_ard_regression()
    N = DictToArrayBijection.map(model.initial_point()).data.shape[0]

    sample_logp_fn = make_pathfinder_sample_fn(
        model,
        N,
        MAXCOR,
        jacobian=True,
        compile_kwargs=COMPILE_KWARGS,
        vectorize=False,
        return_inv_hessian_diag=True,
    )

    cb, rng, _ = _run_streaming_lbfgs(model, sample_logp_fn, seed=20260301)

    assert cb.any_valid
    assert cb.best_state is not None

    state = cb.best_state
    _, _, _, diag_est = sample_logp_fn(
        state["x"],
        state["g"],
        state["alpha"],
        state["s_win"],
        state["z_win"],
        rng.standard_normal((1, N)),
    )
    diag_est = np.asarray(diag_est)
    assert np.all(np.isfinite(diag_est))
    assert np.all(diag_est > 0)

    u = rng.standard_normal((4000, N))
    phi, _, _, _ = sample_logp_fn(
        state["x"],
        state["g"],
        state["alpha"],
        state["s_win"],
        state["z_win"],
        u,
    )
    emp_var = np.var(np.asarray(phi), axis=0, ddof=1)
    assert np.all(np.isfinite(emp_var))
    assert np.all(emp_var > 0)

    np.testing.assert_allclose(emp_var, diag_est, rtol=0.2, atol=1e-3)


@pytest.mark.parametrize("model_name", ["ard_regression"])
@pytest.mark.parametrize("vectorize", [False, True])
def test_rng_reproducibility(model_name, vectorize):
    """Two streaming runs with the same seed produce bit-identical results."""
    model = MODEL_FACTORIES[model_name]()
    fn = _make_single_fn(model, vectorize=vectorize)

    r1 = fn(123)
    r2 = fn(123)

    assert r1.path_status not in FAILED_PATH_STATUS
    assert r2.path_status not in FAILED_PATH_STATUS

    np.testing.assert_array_equal(r1.samples, r2.samples)
    np.testing.assert_array_equal(r1.logP, r2.logP)
    np.testing.assert_array_equal(r1.logQ, r2.logQ)
    assert r1.best_step_idx == r2.best_step_idx if hasattr(r1, "best_step_idx") else True


@pytest.mark.parametrize("model_name", ["ard_regression"])
@pytest.mark.parametrize("vectorize", [False, True])
def test_fixture_match(model_name, vectorize):
    """Streaming path selects an ELBO argmax consistent with the fixture.

    LBFGS is mocked so the streaming path processes the exact same trajectory
    stored in the fixture. ELBO argmax must fall within a generous window.
    """
    x_full, _, elbo_ref = _load_fixture(model_name)
    ref_argmax = int(np.nanargmax(elbo_ref))
    L = x_full.shape[0] - 1

    model = MODEL_FACTORIES[model_name]()
    N = DictToArrayBijection.map(model.initial_point()).data.shape[0]

    with _make_lbfgs_mock(x_full):
        fn = _make_single_fn(model, vectorize=vectorize)
        result = fn(42)

    tag = f"[{model_name}]"

    assert result.path_status not in FAILED_PATH_STATUS, f"{tag} path failed: {result.path_status}"
    assert result.samples is not None
    assert result.samples.shape == (1, NUM_DRAWS, N)
    assert np.any(np.isfinite(result.logP)), f"{tag} all logP are non-finite"
    assert np.any(np.isfinite(result.logQ)), f"{tag} all logQ are non-finite"

    argmax = int(result.elbo_argmax)
    tolerance = max(3, L // 4)
    assert abs(argmax - ref_argmax) <= tolerance, (
        f"{tag} elbo_argmax={argmax} too far from fixture ref_argmax={ref_argmax} "
        f"(L={L}, tolerance={tolerance})"
    )


@pytest.mark.parametrize("model_name", ALL_MODEL_NAMES)
def test_elbo_all_finite(model_name):
    """Reference ELBO values recorded from current code must all be finite."""
    *_, elbo_ref = _load_fixture(model_name)
    assert np.all(np.isfinite(elbo_ref)), (
        f"[{model_name}] reference ELBO contains non-finite values: "
        f"{elbo_ref[~np.isfinite(elbo_ref)]}"
    )


def _check_statistical_equivalence(
    model_name,
    elbo,
    elbo_ref,
    min_corr: float = 0.85,
    min_ballpark_fraction: float = 0.8,
    max_rel_diff: float = 0.1,
):
    """Shared helper: assert that elbo is statistically consistent with elbo_ref.

    All checks must pass; no platform-specific randomness guarantees.
    """
    from scipy.stats import spearmanr

    L = len(elbo_ref)
    if L < 2:
        pytest.skip(f"[{model_name}] only {L} ELBO step(s), statistical check not meaningful")

    assert np.all(np.isfinite(elbo)), f"[{model_name}] ELBO contains non-finite values"

    corr, _ = spearmanr(elbo, elbo_ref)
    assert corr > min_corr, (
        f"[{model_name}] Spearman rank correlation {corr:.3f} < {min_corr}; "
        "ELBO ordering is inconsistent with reference"
    )

    argmax_new = int(np.argmax(elbo))
    argmax_ref = int(np.nanargmax(elbo_ref))
    tol = max(5, L // 5)
    assert abs(argmax_new - argmax_ref) <= tol, (
        f"[{model_name}] argmax {argmax_new} is more than {tol} steps from "
        f"reference argmax {argmax_ref} (L={L})"
    )

    # Ballpark: fraction of values within max_rel_diff of reference
    denom = np.abs(elbo_ref) + 1e-10
    rel_diff = np.abs(elbo - elbo_ref) / denom
    within = np.isfinite(rel_diff) & (rel_diff <= max_rel_diff)
    fraction = np.mean(within)
    pairs = "\n".join(
        f"  step {i}: elbo={e:.6g}, elbo_ref={r:.6g}, rel_diff={d:.4f}"
        for i, (e, r, d) in enumerate(zip(elbo, elbo_ref, rel_diff))
    )
    assert fraction >= min_ballpark_fraction, (
        f"[{model_name}] only {fraction:.1%} of ELBO values within {max_rel_diff:.0%} of "
        f"reference (required {min_ballpark_fraction:.0%})\n"
        f"Full trace (elbo, elbo_ref) pairs:\n{pairs}"
    )


@pytest.mark.parametrize("model_name", ALL_MODEL_NAMES)
@pytest.mark.parametrize("vectorize", [False, True])
def test_elbo_statistical_equivalence(model_name, vectorize):
    """Full-stack: LBFGSStreamingCallback computes alpha/s/z internally; ELBO matches reference."""
    x_full, g_full, elbo_ref = _load_fixture(model_name)
    model = MODEL_FACTORIES[model_name]()
    elbo = _run_callback_replay(model, x_full, vectorize=vectorize)
    assert (
        elbo.shape == elbo_ref.shape
    ), f"[{model_name}] shape mismatch: {elbo.shape} vs {elbo_ref.shape}"
    _check_statistical_equivalence(model_name, elbo, elbo_ref)


@pytest.mark.parametrize("model_name", ["ard_regression"])
@pytest.mark.parametrize("vectorize", [False])
def test_elbo_statistical_equivalence_different_seed(model_name, vectorize):
    """Verify tolerances are meaningful by using a different ELBO seed."""
    x_full, _, elbo_ref = _load_fixture(model_name)
    model = MODEL_FACTORIES[model_name]()
    elbo = _run_callback_replay(model, x_full, seed=ELBO_SEED + 9999, vectorize=vectorize)
    assert elbo.shape == elbo_ref.shape
    _check_statistical_equivalence(model_name, elbo, elbo_ref)
