import itertools, json
import numpy as np

SEED = 7


def subset_stats(theta, m):
    theta = np.asarray(theta, dtype=float)
    d = len(theta)
    subs = np.array(list(itertools.combinations(range(d), m)), dtype=int)
    scores = theta[subs].sum(axis=1)
    scores -= scores.max()
    p = np.exp(scores); p /= p.sum()
    X = np.zeros((len(subs), d))
    X[np.arange(len(subs))[:, None], subs] = 1.0
    mu = p @ X
    cov = (X * p[:, None]).T @ X - np.outer(mu, mu)
    return mu, cov, subs, p


def resistance(cov, i, j):
    pinv = np.linalg.pinv(cov, rcond=1e-11)
    z = np.zeros(cov.shape[0]); z[i] = 1; z[j] = -1
    return float(z @ pinv @ z)


def bound_sweep():
    rows=[]
    for d,m in [(6,2),(8,4),(12,2),(12,4),(12,6)]:
        worst=0.; worst_pair=None; nullerr=0.; min_psd=1e9
        for a in np.linspace(0, 12, 13):
            theta=np.zeros(d); theta[0]=a; theta[1]=-a/2
            mu,cov,_,_=subset_stats(theta,m)
            nullerr=max(nullerr, float(np.max(np.abs(cov@np.ones(d)))))
            v=np.diag(cov); V=v.sum(); B=.5*(np.diag(v)-np.outer(v,v)/V)
            min_psd=min(min_psd, float(np.linalg.eigvalsh(cov-B).min()))
            for i in range(d):
                for j in range(i+1,d):
                    r=resistance(cov,i,j)/(1/v[i]+1/v[j])
                    if r>worst: worst=float(r); worst_pair=(a,i,j)
        rows.append(dict(d=d,m=m,max_resistance_ratio=worst,worst_case=worst_pair,
                         max_null_residual=nullerr,min_cov_minus_bound_eigenvalue=min_psd))
    return rows


def trust_scaling():
    d,m=12,4; theta=np.linspace(2,-2,d); mu,cov,_,_=subset_stats(theta,m)
    g=np.linspace(-1,1,d); g-=g.mean()
    u=np.linalg.pinv(cov, rcond=1e-11)@g; u-=u.mean()
    v=np.diag(cov)
    raw=max(abs(u[i]-u[j])/np.sqrt(1/v[i]+1/v[j]) for i in range(d) for j in range(i+1,d))
    out=[]
    for rho in [0.01,0.03,0.1,0.3,1.0]:
        scale=min(1.,rho/raw); delta=-u*scale
        observed=max(abs(delta[i]-delta[j])/np.sqrt(1/v[i]+1/v[j]) for i in range(d) for j in range(i+1,d))
        out.append(dict(rho=rho,raw_normalized_step=raw,observed=observed,predicted=min(rho,raw)))
    return out


def optimize_compare():
    d,m=12,4
    target,_c,_,_=subset_stats(np.array([3.,2.,1.,.5,0,0,0,0,0,0,-1.,-2.]),m)
    results=[]
    for method in ['vanilla','natural_trust']:
        theta=np.zeros(d); losses=[]; cvs=[]
        for t in range(80):
            mu,cov,_,_=subset_stats(theta,m)
            err=mu-target; loss=.5*float(err@err)
            losses.append(loss); cvs.append(float(np.std(mu)/np.mean(mu)))
            g=cov@err
            if method=='vanilla':
                delta=-1.4*g
            else:
                u=np.linalg.pinv(cov,rcond=1e-11)@g; u-=u.mean()
                v=np.diag(cov); rho=.35
                mx=max(abs(u[i]-u[j])/np.sqrt(1/v[i]+1/v[j]) for i in range(d) for j in range(i+1,d))
                delta=-1.4*u*min(1.,rho/mx) if mx>0 else np.zeros(d)
            theta += delta
        results.append(dict(method=method,initial_loss=losses[0],final_loss=losses[-1],
                            loss_step_10=losses[10],loss_step_40=losses[40],
                            final_load_cv=cvs[-1],min_loss=float(min(losses))))
    return results


def main():
    report={'seed':SEED,
      'predictions':{
        'resistance_bound':'max normalized resistance ratio <= 1 for every field and pair',
        'covariance_lower_bound':'Sigma - 0.5(D-vv^T/V) is PSD',
        'trust_scaling':'normalized max pairwise update equals min(rho, raw normalized step)'},
      'bound_sweep':bound_sweep(),'trust_sweep':trust_scaling(),
      'optimization':optimize_compare()}
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()