"""Smoke tests for factored_policy core model (fast, local)."""
import numpy as np
import sys, os
sys.path.insert(0, os.path.dirname(__file__))
import factored_policy as F


def test_Q_matches_montecarlo():
    levels = [3, 4, 2]
    beta, gamma = F.random_model(levels, seed=42, gamma_scale=0.7)
    pi_a = F.free_to_policy(np.array([0.2, 0.3, 0.1, 0.2, 0.3, 0.4]), levels)
    pi_b = F.free_to_policy(np.array([0.25, 0.25, 0.2, 0.1, 0.4, 0.3]), levels)
    Q_closed = F.Q_value(beta, gamma, levels, pi_a, pi_b, mu=0.5)
    # Monte Carlo: sample profiles, average the linear-predictor win prob
    rng = np.random.default_rng(1); N = 400_000
    ta = np.stack([rng.choice(levels[d], size=N, p=pi_a[d]) for d in range(len(levels))], axis=1)
    tb = np.stack([rng.choice(levels[d], size=N, p=pi_b[d]) for d in range(len(levels))], axis=1)
    g_a = np.zeros(N); g_b = np.zeros(N)
    for d in range(len(levels)):
        for l in range(levels[d]-1):
            g_a[ta[:, d] == l] += beta[d][l]
            g_b[tb[:, d] == l] += beta[d][l]
    for (d1, d2), G in gamma.items():
        for i in range(G.shape[0]):
            for j in range(G.shape[1]):
                g_a[(ta[:, d1] == i) & (ta[:, d2] == j)] += G[i, j]
                g_b[(tb[:, d1] == i) & (tb[:, d2] == j)] += G[i, j]
    Q_mc = 0.5 + np.mean(g_a - g_b)
    assert abs(Q_closed - Q_mc) < 0.01, (Q_closed, Q_mc)
    print(f"[ok] Q closed={Q_closed:.4f} MC={Q_mc:.4f}")


def test_closed_form_is_optimum():
    levels = [3, 4, 3, 2]
    beta, gamma = F.random_model(levels, seed=7, gamma_scale=0.4)
    n = F.free_length(levels)
    p_free = np.concatenate([np.full(L-1, 1.0/L) for L in levels])
    lam = 3.0
    x_star = F.closed_form_policy(beta, gamma, levels, lam, p_free)
    pi_star = F.free_to_policy(x_star, levels)
    p_ref = F.uniform_policy(levels)
    def obj(x):
        pi = F.free_to_policy(x, levels)
        full = np.concatenate([pi[d] for d in range(len(levels))])
        full_p = np.concatenate([p_ref[d] for d in range(len(levels))])
        return F.Q_value(beta, gamma, levels, pi, p_ref) - lam*np.sum((full-full_p)**2)
    val_star = obj(x_star)
    # random feasible search
    rng = np.random.default_rng(3); best = -1e9
    for _ in range(20000):
        x = np.zeros(n); i=0
        for L in levels:
            v = rng.dirichlet(np.ones(L)); x[i:i+L-1] = v[:-1]; i += L-1
        best = max(best, obj(x))
    assert val_star >= best - 1e-6, (val_star, best)
    print(f"[ok] closed-form obj={val_star:.5f} >= random best={best:.5f}")


def test_payoff_matrix_is_zero_sum():
    levels = [2, 2, 3]
    beta, gamma = F.random_model(levels, seed=11)
    profiles, W = F.payoff_matrix_full(beta, gamma, levels, mu=0.5)
    assert np.allclose(W + W.T, 1.0, atol=1e-9)
    v_max, _ = F.maximin_value_lp(W)
    v_min, _ = F.minimax_value_lp(W)
    assert abs(v_max - v_min) < 1e-7, (v_max, v_min)
    assert abs(v_max - 0.5) < 1e-6, v_max
    print(f"[ok] antisymmetric W, von Neumann value={v_max:.6f}=0.5")


if __name__ == "__main__":
    test_Q_matches_montecarlo()
    test_closed_form_is_optimum()
    test_payoff_matrix_is_zero_sum()
    print("ALL CORE TESTS PASSED")