import json, math
import numpy as np


def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-np.clip(x, -60, 60)))


def weights(losses, tau):
    z = -(losses - np.min(losses)) / max(tau, 1e-15)
    z -= np.max(z)
    w = np.exp(np.clip(z, -80, 0))
    return w / w.sum()


def entropy(w):
    return float(-(w * np.log(np.maximum(w, 1e-300))).sum())


def controller_rate(h, alpha=.05, eps=.02, hc=.5, delta=.05):
    b = sigmoid((h-hc)/delta)
    return alpha * (eps + (1-eps)*b)


def main():
    rng = np.random.default_rng(1234)
    N = 128; hc=.5; delta=.05; alpha=.05; eps=.02
    # Prediction 1: ESS identity and transition ESS, across random populations.
    identity_errors=[]
    for _ in range(100):
        losses=rng.normal(size=N)*rng.uniform(.1, 5)
        w=weights(losses, rng.uniform(.01, 3))
        H=entropy(w); h=H/math.log(N); ess=math.exp(H)
        identity_errors.append(abs(ess/N - math.exp(-(1-h)*math.log(N))))
    transition_pred=N**hc
    # Make a population with exactly uniform weights over K candidates; h=log(K)/log(N).
    K=round(transition_pred)
    w=np.zeros(N); w[:K]=1/K
    h_transition=entropy(w)/math.log(N)
    ess_transition=math.exp(entropy(w))

    # Prediction 2: tau log slope is controller_rate in fixed-entropy regimes.
    rate_rows=[]
    for h in [.1, .9, .5]:
        tau=1.0; logs=[]
        for _ in range(200):
            logs.append(math.log(tau)); tau *= math.exp(-controller_rate(h,alpha,eps,hc,delta))
        observed=-(logs[-1]-logs[0])/(len(logs)-1)
        rate_rows.append({'h':h, 'predicted_abs_slope':controller_rate(h,alpha,eps,hc,delta), 'observed_abs_slope':observed})

    # Prediction 3: entropy threshold crossing in a temperature sweep. Candidate losses are
    # 0 for K good candidates and 1 for the rest; solve/search tau where h ~= hc.
    sweep=[]
    for K0 in [1, 2, 4, 8, 16, 32, 64]:
        taus=np.logspace(-3,2,500)
        hs=[]
        for tau in taus:
            ll=np.ones(N); ll[:K0]=0
            hs.append(entropy(weights(ll,tau))/math.log(N))
        idx=int(np.argmin(np.abs(np.asarray(hs)-hc)))
        sweep.append({'good_candidates':K0, 'tau_at_hc':float(taus[idx]), 'h':float(hs[idx])})

    # Mini optimizer: diagonal ill-conditioned quadratic, same perturbations for all methods per seed.
    d=24; T=180; pop=64; sigma=.18; gamma=1.0
    curv=np.geomspace(1, 20, d)
    def run(seed, mode):
        r=np.random.default_rng(seed); theta=r.normal(0, 2, d); tau=1.0
        vals=[]; hs=[]
        for t in range(T):
            e=r.normal(size=(pop,d)); cand=theta[None,:]+sigma*e
            loss=.5*np.sum(curv[None,:]*cand*cand,axis=1)
            if mode=='constant': tau_use=1.0
            else: tau_use=tau
            w=weights(loss,tau_use); H=entropy(w); h=H/math.log(pop)
            theta=(1-gamma)*theta+gamma*(w[:,None]*cand).sum(axis=0)
            if mode=='feedback': tau*=math.exp(-controller_rate(h,alpha,eps,hc,delta))
            elif mode=='fixed': tau*=math.exp(-alpha)
            vals.append(float(.5*np.sum(curv*theta*theta))); hs.append(h)
        return {'final':vals[-1], 'best':min(vals), 'mean_h_last20':float(np.mean(hs[-20:])), 'tau_final':tau, 'loss_curve':vals}
    results={}
    for mode in ['feedback','fixed','constant']:
        runs=[run(s,mode) for s in range(5)]
        results[mode]={'final_mean':float(np.mean([x['final'] for x in runs])), 'final_std':float(np.std([x['final'] for x in runs])), 'best_mean':float(np.mean([x['best'] for x in runs])), 'h_last20_mean':float(np.mean([x['mean_h_last20'] for x in runs])), 'tau_final_mean':float(np.mean([x['tau_final'] for x in runs]))}

    out={'config':{'N':N,'hc':hc,'delta':delta,'alpha':alpha,'epsilon':eps},
         'prediction_1_ess_identity_max_abs_error':max(identity_errors),
         'prediction_1_transition':{'predicted_ESS':transition_pred,'constructed_K':K,'observed_h':h_transition,'observed_ESS':ess_transition},
         'prediction_2_rates':rate_rows,
         'prediction_3_temperature_sweep':sweep,
         'mini_experiment':results}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()