import heapq
import json
import time
import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

SEED = 7

class EventGraph:
    def __init__(self, n_source, n_hidden, n_out, source_edges, hidden_edges, vmax=100):
        self.S, self.H, self.O = n_source, n_hidden, n_out
        self.source_edges = source_edges
        self.hidden_edges = hidden_edges
        self.vmax = vmax

    def run(self, x, return_trace=False):
        a = np.zeros(self.H, dtype=float)
        visits = np.zeros(self.H, dtype=int)
        z = np.zeros(self.O, dtype=float)
        q, serial = [], 0
        for i, value in enumerate(x):
            heapq.heappush(q, (0, serial, i, float(value)))
            serial += 1
        events = 0
        trace = []
        while q and events < 100000:
            t, _, node, value = heapq.heappop(q)
            events += 1
            if node < self.S:
                for si, h, w, d in self.source_edges:
                    if si == node:
                        heapq.heappush(q, (t + d, serial, self.S + h, w * value))
                        serial += 1
                continue
            h = node - self.S
            if h >= self.H:
                z[h - self.H] += value
                continue
            if visits[h] >= self.vmax:
                continue
            a[h] += value
            visits[h] += 1
            y = np.tanh(a[h])
            trace.append((t, h, a[h], y, visits[h]))
            for hh, out, w, d in self.hidden_edges:
                if hh == h:
                    heapq.heappush(q, (t + d, serial, self.S + self.H + out, w * y))
                    serial += 1
        result = dict(z=z, accum=a, visits=visits, events=events, trace=trace)
        return result if return_trace else z


def prediction_sweeps():
    rng = np.random.default_rng(SEED)
    rows = []
    # Prediction 1: with one shared node and linear response, final accumulator is sum of arrivals.
    # For k equal arrivals c, a = k*c exactly; measured slope should be 1.
    ks = np.arange(1, 9)
    measured = []
    for k in ks:
        edges = [(i, 0, 1.0, int(i % 2)) for i in range(k)]
        g = EventGraph(k, 1, 1, edges, [(0, 0, 1.0, 0)])
        measured.append(g.run(np.full(k, 0.1), True)['accum'][0])
    slope = float(np.polyfit(ks, measured, 1)[0] / 0.1)
    rows.append({'prediction': 'shared accumulator slope versus number of equal arrivals = 1/cancel-free sum',
                  'predicted': 1.0, 'observed': slope, 'relative_error': abs(slope-1.0)})

    # Prediction 2: tanh response saturates; equal arrivals k*c approach 1, with
    # inverse-tanh(0.9)/c predicted for the 0.9 crossing.
    c = 0.2
    crossing = None
    ys = []
    for k in range(1, 31):
        edges = [(i, 0, 1.0, 0) for i in range(k)]
        g = EventGraph(k, 1, 1, edges, [(0, 0, 1.0, 0)])
        y = g.run(np.full(k, c), True)['trace'][-1][3]
        ys.append(y)
        if crossing is None and y >= .9:
            crossing = k
    predicted_crossing = int(np.ceil(np.arctanh(.9) / c))
    rows.append({'prediction': 'shared tanh response reaches 0.9 at k=ceil(atanh(0.9)/c)',
                  'predicted': predicted_crossing, 'observed': crossing,
                  'relative_error': abs(crossing-predicted_crossing)/predicted_crossing})

    # Prediction 3: visit cap Vmax truncates exactly after Vmax hidden visits.
    cap_rows = []
    k = 8
    for cap in range(1, 6):
        edges = [(i, 0, 1.0, 0) for i in range(k)]
        g = EventGraph(k, 1, 1, edges, [(0, 0, 1.0, 0)], vmax=cap)
        r = g.run(np.ones(k), True)
        cap_rows.append((cap, int(r['visits'][0]), float(r['accum'][0])))
    rows.append({'prediction': 'visit cap produces exactly min(k,Vmax) accepted updates',
                 'predicted': [min(k, c) for c in range(1, 6)],
                 'observed': [x[1] for x in cap_rows],
                 'relative_error': 0.0 if all(x[1] == min(k, x[0]) for x in cap_rows) else 1.0})
    return rows


def make_features(X, shared=True, P=None, Q=None, seed=SEED):
    rng = np.random.default_rng(seed)
    n, d = X.shape
    S, H, O = 16, 24, 10
    if P is None:
        P = rng.normal(0, 1/np.sqrt(d), (d, S))
    if Q is None:
        Q = rng.normal(0, 1/np.sqrt(H), (H, O))
    out = []
    event_counts = []
    for x in X:
        src = x @ P
        if shared:
            edges = [(i, int(i % 8), 0.55, i % 3) for i in range(S)]
        else:
            edges = [(i, i, 0.55, 0) for i in range(S)]
        hidden_edges = [(h, o, Q[h, o], 0) for h in range(H) for o in range(O)]
        g = EventGraph(S, H, O, edges, hidden_edges, vmax=20)
        r = g.run(src, return_trace=True)
        out.append(r['z'])
        event_counts.append(r['events'])
    return np.asarray(out), float(np.mean(event_counts))


def classification_test():
    digits = load_digits()
    X = digits.data.astype(float) / 16.0
    y = digits.target
    Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=.25, random_state=SEED, stratify=y)
    rng = np.random.default_rng(SEED)
    P = rng.normal(0, 1/np.sqrt(X.shape[1]), (X.shape[1], 16))
    Q = rng.normal(0, 1/np.sqrt(24), (24, 10))
    result = {}
    for shared in (False, True):
        t0 = time.perf_counter()
        Ftr, etr = make_features(Xtr, shared=shared, P=P, Q=Q)
        Fte, ete = make_features(Xte, shared=shared, P=P, Q=Q)
        clf = LogisticRegression(max_iter=300, C=1.0, random_state=SEED)
        clf.fit(Ftr, ytr)
        pred = clf.predict(Fte)
        result['shared' if shared else 'no_sharing'] = {
            'accuracy': float(accuracy_score(yte, pred)),
            'seconds': time.perf_counter() - t0,
            'train_mean_events': etr, 'test_mean_events': ete,
            'feature_dim': int(Ftr.shape[1])
        }
    return result


def main():
    report = {'seed': SEED, 'mechanism_checks': prediction_sweeps(), 'classification': classification_test()}
    with open('results.json', 'w') as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()