import numpy as np
from router import exact_stats, sample_subset_dp, natural_direction

theta=np.array([2.0,1.0,0.2,-0.4,-1.0,-2.0])
m=2
mu,sigma=exact_stats(theta,m)
rng=np.random.default_rng(11)
counts=np.zeros(len(theta))
N=200000
for _ in range(N):
    counts[sample_subset_dp(theta,m,rng)] += 1
emp=counts/N
err=np.max(np.abs(emp-mu))
g=np.array([1.2,-.7,.4,.1,-.5,-.5])
d,info=natural_direction(theta,m,g,rho=.17)
v=np.diag(info['sigma'])
cap=max(abs(d[i]-d[j])/np.sqrt(1/v[i]+1/v[j]) for i in range(len(theta)) for j in range(i+1,len(theta)))
print({'exact_mu':mu.tolist(),'empirical_mu':emp.tolist(),'max_sampling_error':float(err), 'sum_direction':float(d.sum()), 'cap':float(cap),'rho':.17,'raw_cap':info['raw_cap']})
assert err < .006
assert abs(d.sum()) < 1e-10
assert cap <= .17 + 1e-10