import json
import numpy as np
from scipy.linalg import eigh
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

EPS = 1e-10

def sym(a): return (a + a.T) / 2

def spectral(a, fn):
    v, u = eigh(sym(a)); v = np.maximum(v, EPS)
    return sym((u * fn(v)) @ u.T)

def sqrt_spd(a): return spectral(a, np.sqrt)
def invsqrt_spd(a): return spectral(a, lambda x: 1 / np.sqrt(x))
def spd_log(a): return spectral(a, np.log)
def spd_exp(a): return spectral(a, np.exp)
def project(a, lo, hi): return spectral(a, lambda x: np.clip(x, lo, hi))

def bw2(p, q):
    hp = sqrt_spd(p); r = sqrt_spd(hp @ q @ hp)
    return float(max(np.trace(p + q) - 2 * np.trace(r), 0))

def objective(s, mats, weights):
    return .5 * sum(w * bw2(s, a) for w, a in zip(weights, mats))

def eigvals(a): return eigh(sym(a), eigvals_only=True)
def condition(a):
    v = eigvals(a); return float(v[-1] / max(v[0], EPS))

def bw_update(s, mats, weights, bounds=None):
    hs = sqrt_spd(s); z = np.zeros_like(s)
    for w, a in zip(weights, mats): z += w * sqrt_spd(hs @ a @ hs)
    his = invsqrt_spd(s)
    x = his @ (z @ z) @ his                 # article.md Eq. (3)
    return project(x, *bounds) if bounds else sym(x)

def arithmetic(mats, w): return sum(x * a for x, a in zip(w, mats))
def log_euclidean(mats, w): return spd_exp(sum(x * spd_log(a) for x, a in zip(w, mats)))

def random_spd(rng, d, lo, hi):
    q, _ = np.linalg.qr(rng.normal(size=(d, d)))
    v = np.exp(rng.uniform(np.log(lo), np.log(hi), d))
    return sym((q * v) @ q.T)

def math_check(rng):
    lo, hi, d = .01, 3., 2
    # Direct numerical check of the claimed non-expansiveness.
    ratios = []
    for _ in range(120):
        x, y = random_spd(rng, 5, .01, 20), random_spd(rng, 5, .01, 20)
        ratios.append(np.sqrt(bw2(project(x,.1,10), project(y,.1,10))) / max(np.sqrt(bw2(x,y)),1e-12))
    # Search the paper's stated type of transient floor exit: feasible inputs and S0,
    # but an unprojected unit step below lo. This is a diagnostic, not cherry-picked output.
    found = None
    for trial in range(1200):
        mats = [random_spd(rng, d, lo, hi) for _ in range(3)]
        w = np.ones(3) / 3; s = arithmetic(mats, w)
        if eigvals(s)[0] < lo - 1e-8: continue
        raw, clipped = [s], [s]
        for _ in range(5):
            raw.append(bw_update(raw[-1], mats, w))
            clipped.append(bw_update(clipped[-1], mats, w, (lo, hi)))
        if min(eigvals(x)[0] for x in raw) < lo - 1e-6:
            found = (mats, w, raw, clipped, trial + 1); break
    if found is None:
        # Still provide a reproducible ordinary feasible run.
        mats = [random_spd(rng, d, lo, hi) for _ in range(3)]; w=np.ones(3)/3
        s=arithmetic(mats,w); raw=[s]; clipped=[s]
        for _ in range(5): raw.append(bw_update(raw[-1],mats,w)); clipped.append(bw_update(clipped[-1],mats,w,(lo,hi)))
        trials=None
    else: mats,w,raw,clipped,trials=found
    return {
      'projection_max_ratio':float(max(ratios)), 'projection_mean_ratio':float(np.mean(ratios)),
      'exit_search_trials':trials, 'raw_min_eigenvalues':[float(eigvals(x)[0]) for x in raw],
      'projected_min_eigenvalues':[float(eigvals(x)[0]) for x in clipped],
      'raw_objective':[objective(x,mats,w) for x in raw],
      'projected_objective':[objective(x,mats,w) for x in clipped],
      'raw_condition':[condition(x) for x in raw], 'projected_condition':[condition(x) for x in clipped]
    }

def pooling_task(rng):
    n,d,m=500,5,6; q,_=np.linalg.qr(rng.normal(size=(d,d)))
    bases=[q@np.diag([4,2,1,.5,.25])@q.T,q@np.diag([.25,.5,1,2,4])@q.T]
    feats={k:[] for k in ('arith','log','bw','pbw')}; mats_by={k:[] for k in feats}; all_inputs=[]; labels=[]
    for j in range(n):
        base=bases[j%2]; mats=[]
        for _ in range(m):
            z=rng.normal(size=(d,d)); mats.append(sym(base+.15*(z@z.T/d)+.03*np.eye(d)))
        w=np.ones(m)/m; ar=arithmetic(mats,w); vals={'arith':ar,'log':log_euclidean(mats,w)}
        bw=ar; pbw=ar
        for _ in range(4): bw=bw_update(bw,mats,w); pbw=bw_update(pbw,mats,w,(.2,5))
        vals.update({'bw':bw,'pbw':pbw}); all_inputs.append(mats); labels.append(j%2)
        for k,s in vals.items(): mats_by[k].append(s); feats[k].append(s[np.triu_indices(d)])
    labels=np.asarray(labels); tr,te=train_test_split(np.arange(n),test_size=.3,random_state=7,stratify=labels); out={}
    for k in feats:
        x=np.asarray(feats[k]); clf=LogisticRegression(max_iter=1000,random_state=7).fit(x[tr],labels[tr])
        out[k]={'accuracy':float(accuracy_score(labels[te],clf.predict(x[te]))), 'mean_condition':float(np.mean([condition(a) for a in mats_by[k]])), 'mean_objective':float(np.mean([objective(s,aa,np.ones(m)/m) for s,aa in zip(mats_by[k],all_inputs)]))}
    return out

def main():
    result={'math_check':math_check(np.random.default_rng(3150)),'pooling_task':pooling_task(np.random.default_rng(3151))}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()