import json
import numpy as np
from diversity_baseline import diversity_baseline, uniform_loo, effective_count


def main():
    out = {}
    # Prediction 1: p=0 is exactly uniform LOO, independent of embeddings.
    rng = np.random.default_rng(2183)
    costs = rng.normal(size=6)
    emb = rng.normal(size=(6, 3))
    b0, a0, q0, _ = diversity_baseline(costs, emb, p=0)
    bu = uniform_loo(costs)
    out['p0_max_baseline_error'] = float(np.max(np.abs(b0 - bu)))
    out['p0_effective_count_observed_predicted'] = [float(effective_count(q0).mean()), 5.0]

    # Prediction 2: for two peer distances, weight ratio is (d_far/d_near)^p.
    # eps is negligible here; use a group with one anchor and two peers.
    costs2 = np.array([0., 1., 2.])
    emb2 = np.array([[0., 0.], [1., 0.], [4., 0.]])
    _, _, q2, d2 = diversity_baseline(costs2, emb2, p=1, eps=1e-8)
    observed_ratio = q2[0, 2] / q2[0, 1]
    predicted_ratio = (d2[0, 2] + 1e-8) / (d2[0, 1] + 1e-8)
    out['distance_ratio_observed_predicted'] = [float(observed_ratio), float(predicted_ratio)]

    # Prediction 3: duplicated/near peers are downweighted; sweep p and compare
    # effective count with the exact two-distance prediction for the anchor.
    # Anchor at 0, two near peers at +/-delta and two far peers at +/-D.
    delta, D = 0.1, 2.0
    emb3 = np.array([[0., 0.], [delta, 0.], [-delta, 0.], [D, 0.], [-D, 0.]])
    costs3 = np.arange(5, dtype=float)
    sweep = []
    for p in [0., .5, 1., 2.]:
        _, _, q, d = diversity_baseline(costs3, emb3, p=p)
        near = (d[0, 1] + 1e-8) ** p
        far = (d[0, 3] + 1e-8) ** p
        predicted_eff = 1.0 / (2 * (near/(2*near+2*far))**2 + 2 * (far/(2*near+2*far))**2)
        sweep.append({'p': p, 'effective_observed': float(effective_count(q)[0]),
                      'effective_predicted': float(predicted_eff),
                      'far_to_near_weight_observed': float(q[0,3]/q[0,1]),
                      'far_to_near_weight_predicted': float(far/near)})
    out['effective_count_sweep'] = sweep

    # Small grouped-policy proxy: categorical logits, score-function gradients.
    # Actions 0/1 are near-duplicates and cheap; 2/3 are structurally distinct
    # and expensive. Diversity weighting should alter, not automatically improve,
    # gradient variance; this is an honest secondary comparison.
    probs = np.array([.35, .35, .15, .15])
    action_emb = np.array([[0.,0.], [.02,0.], [1.,0.], [1.02,0.]])
    action_cost = np.array([1., 1.1, 4., 4.1])
    n_groups, B = 20000, 8
    grads_u, grads_d = [], []
    for _ in range(n_groups):
        acts = rng.choice(4, size=B, p=probs)
        c = action_cost[acts]
        z = action_emb[acts]
        bu = uniform_loo(c)
        bd, ad, _, _ = diversity_baseline(c, z, p=1.)
        au = bu - c
        score = np.eye(4)[acts] - probs
        grads_u.append(((-au[:,None] * score)).mean(0))
        grads_d.append(((-ad[:,None] * score)).mean(0))
    var_u = np.var(grads_u, axis=0, ddof=1).mean()
    var_d = np.var(grads_d, axis=0, ddof=1).mean()
    out['gradient_variance_uniform_diversity_ratio'] = [float(var_u), float(var_d), float(var_d/var_u)]
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()