import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import argparse, json, numpy as np
from lamm.state import AssociativeMemory

def run(seed=314159, d=48, lam=.001, steps=200):
    rng=np.random.default_rng(seed)
    m=AssociativeMemory(d,lam)
    e=[]
    for _ in range(steps):
        m.update(rng.normal(size=d)*.01,rng.normal(size=d)*.01)
        e.append(m.energy())
    return {
        "lambda":lam,
        "mean_energy":float(np.mean(e)),
        "std_energy":float(np.std(e)),
        "min_energy":float(np.min(e)),
        "max_energy":float(np.max(e)),
        "final_energy":float(e[-1])
    }

for lam in [.0001,.001,.01,.1]:
    print(json.dumps(run(lam=lam)))