import json, math, random
import numpy as np
from sklearn.datasets import load_digits
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

SEED = 1248
np.random.seed(SEED)
random.seed(SEED)


def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-np.clip(z, -40, 40)))


def math_checks():
    out = {}
    # Prediction 1: around target rho, homeostatic scalar update contracts with
    # q = 1 - eta*lambda*[rho(1-rho)]^2.
    rho, eta, z0, tol = .7, .8, -2., 1e-3
    alpha_r = rho * (1-rho)
    rows = []
    for lam in [.05, .1, .2, .4, .8]:
        z = z0
        q = 1 - eta * lam * alpha_r**2
        pred = math.log(tol / abs(sigmoid(z0)-rho)) / math.log(abs(q))
        observed = None
        for k in range(200000):
            h = sigmoid(z)
            z += eta * lam * h * (1-h) * (rho-h)
            if abs(sigmoid(z)-rho) < tol:
                observed = k + 1
                break
        rows.append({'lambda': lam, 'observed_steps': observed,
                     'predicted_linear_steps': pred, 'predicted_q': q})
    out['homeostatic_contraction'] = rows

    # Prediction 2: with scalar A=[a], C_A=a^2, so temporal drive scales as a^2.
    dh = .23
    rows = []
    for a in [.25, .5, 1., 2., 4.]:
        rows.append({'a': a, 'observed_ratio': (a*a*dh)/dh,
                     'predicted_ratio': a*a})
    out['projection_scaling'] = rows

    # Prediction 3: alpha is maximal at h=.5 and goes to zero at saturation.
    hs = [.01, .05, .1, .25, .5, .75, .9, .95, .99]
    out['gain_suppression'] = [
        {'h': h, 'alpha': h*(1-h), 'predicted_relative': 4*h*(1-h)}
        for h in hs]
    return out


def local_train(X, C=24, M=12, epochs=3, eta=.22, etab=.01,
                lam=.12, rho=.5, ordered=True, seed=0):
    rng = np.random.RandomState(seed)
    n = X.shape[1]
    W = rng.normal(0, .18, (C, n))
    b = np.zeros(C)
    A = rng.normal(0, 1/math.sqrt(M), (M, C))
    order = np.arange(len(X))
    for _ in range(epochs):
        if not ordered:
            rng.shuffle(order)
        for p in range(len(order)-1):
            i, j = order[p], order[p+1]
            x, xn = X[i], X[j]
            h = sigmoid(W @ x + b)
            hn = sigmoid(W @ xn + b)
            u = A.T @ (A @ (hn-h)) - lam*(h-rho)
            gain = h*(1-h)
            d = gain*u
            W += eta*np.outer(d, x)
            b += etab*d
    return W, b


def bp_train(X, C=24, epochs=3, lr=.08, seed=0):
    # Backprop baseline: minimize consecutive representation mismatch plus
    # weak variance/homeostatic regularization, using the same ordered stream.
    rng = np.random.RandomState(seed)
    n = X.shape[1]
    W = rng.normal(0, .18, (C, n)); b = np.zeros(C)
    for _ in range(epochs):
        for i in range(len(X)-1):
            x, xn = X[i], X[i+1]
            z=W@x+b; zn=W@xn+b
            h=sigmoid(z); hn=sigmoid(zn)
            diff=hn-h
            # loss .5||hn-h||^2 + .08 mean((h-.5)^2), update both pair terms
            gh = -diff + .16*(h-.5)/C
            ghn = diff + .16*(hn-.5)/C
            dz=gh*h*(1-h); dzn=ghn*hn*(1-hn)
            W -= lr*(np.outer(dz,x)+np.outer(dzn,xn))
            b -= lr*(dz+dzn)
    return W,b


def features(X, W, b):
    return sigmoid(X @ W.T + b)


def probe_accuracy(Ftr, ytr, Fte, yte):
    clf = LogisticRegression(max_iter=300, random_state=SEED)
    clf.fit(Ftr, ytr)
    return float(accuracy_score(yte, clf.predict(Fte)))


def experiment():
    d = load_digits()
    X = d.data.astype(np.float64)/16.0
    y = d.target
    # Temporal stream: sort by label, with small within-class noise permutation,
    # making adjacency meaningful but not identical. Hold out a fixed suffix.
    rng=np.random.RandomState(SEED)
    tr, te = train_test_split(np.arange(len(X)), test_size=.3, stratify=y,
                              random_state=SEED)
    tr=np.array(tr); te=np.array(te)
    stream=np.concatenate([tr[np.argsort(y[tr])], te])
    Xtr, ytr = X[stream[:len(tr)]], y[stream[:len(tr)]]
    Xte, yte = X[te], y[te]
    results={}
    configs=[('local_ordered',True),('local_shuffled',False)]
    for name, ordered in configs:
        W,b=local_train(Xtr, ordered=ordered, seed=SEED)
        Ft,Fe=features(Xtr,W,b),features(Xte,W,b)
        results[name]={'accuracy':probe_accuracy(Ft,ytr,Fe,yte),
          'mean_activity':float(Fe.mean()),
          'saturated_fraction':float(np.mean((Fe<.05)|(Fe>.95))),
          'temporal_mse':float(np.mean((Ft[1:]-Ft[:-1])**2))}
    W,b=bp_train(Xtr,seed=SEED)
    Ft,Fe=features(Xtr,W,b),features(Xte,W,b)
    results['backprop_ordered']={'accuracy':probe_accuracy(Ft,ytr,Fe,yte),
      'mean_activity':float(Fe.mean()),
      'saturated_fraction':float(np.mean((Fe<.05)|(Fe>.95))),
      'temporal_mse':float(np.mean((Ft[1:]-Ft[:-1])**2))}
    return results


if __name__ == '__main__':
    result={'math_checks':math_checks(), 'experiment':experiment()}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))