"""Train a tied, bias-free 2 -> 1 -> 2 autoencoder on controlled features.

Only NumPy is needed. All initializations are retained. No test-based selection.
Run: python experiment.py
"""
import argparse
import json
import math
from pathlib import Path
import numpy as np

SEEDS = [7, 11, 23, 42, 101]
PS = [.05, .2, .5]
MODES = ['exclusive', 'independent', 'paired']


def data(p, mode, n, seed):
    rng = np.random.default_rng(seed)
    q = {'exclusive': 0., 'independent': p*p, 'paired': p}[mode]
    # Neither, feature 1 only, feature 2 only, both. Marginals are both p.
    probabilities = [1-2*p+q, p-q, p-q, q]
    masks = np.array([[0., 0.], [1., 0.], [0., 1.], [1., 1.]])
    states = rng.choice(4, size=n, p=probabilities)
    values = masks[states] * rng.uniform(.5, 1.5, size=(n, 2))
    return values, states, probabilities


def forward(x, w, relu):
    h = x @ w
    pre = h[:, None]*w[None, :]
    return np.maximum(pre, 0) if relu else pre


def gradient(x, w, relu):
    h = x @ w
    pre = h[:, None]*w[None, :]
    pred = np.maximum(pre, 0) if relu else pre
    d = (pred-x)/len(x)  # derivative of mean over examples AND two outputs
    if relu:
        d = d*(pre > 0)
    return (d*h[:, None] + (d @ w)[:, None]*x).sum(axis=0)


def checks():
    rng = np.random.default_rng(37)
    x = rng.uniform(.5, 1.5, (31, 2))
    for w in [np.array([.8, -.6]), np.array([.2, .7])]:
        for relu in [False, True]:
            exact = gradient(x, w, relu)
            numerical = []
            for i in range(2):
                delta = np.zeros(2)
                delta[i] = 1e-6
                numerical.append((((forward(x, w+delta, relu)-x)**2).mean()
                                  - ((forward(x, w-delta, relu)-x)**2).mean())/2e-6)
            np.testing.assert_allclose(exact, numerical, rtol=1e-5, atol=1e-8)
    inputs = np.array([[1.,0.], [0.,1.], [0.,0.], [1.,1.]])
    predicted = forward(inputs, np.array([1.,-1.]), True)
    np.testing.assert_allclose(predicted[:3], inputs[:3])
    np.testing.assert_allclose(predicted[3], [0.,0.])
    for p in PS:
        for mode in MODES:
            _, _, probs = data(p, mode, 1, 0)
            assert min(probs) >= 0 and abs(sum(probs)-1) < 1e-12
            assert abs(probs[1]+probs[3]-p) < 1e-12
            assert abs(probs[2]+probs[3]-p) < 1e-12
    print('Checks passed: tied-weight gradients, signed decoder, distribution marginals.', flush=True)


def metrics(x, states, w, relu):
    prediction = forward(x, w, relu)
    mse = float(((prediction-x)**2).mean())
    zero_mse = float((x*x).mean())
    counts = {name: int((states==i).sum()) for i, name in enumerate(['neither','first','second','both'])}
    conditional = {}
    for name, mask in [('single', (states==1)|(states==2)), ('both', states==3)]:
        conditional[name] = float(((prediction[mask]-x[mask])**2).mean()) if mask.any() else None
    return {'mse': mse, 'relative_error': mse/zero_mse,
            'conditional_mse': conditional, 'test_state_counts': counts}


def train(train_x, steps):
    # Paired initial weights and batches for linear/ReLU at each seed.
    initial = np.stack([np.random.default_rng(seed).normal(0, .5, 2) for seed in SEEDS])
    w = np.concatenate([initial, initial]).copy()
    relu = np.array([False]*len(SEEDS)+[True]*len(SEEDS))
    moment = np.zeros_like(w)
    variance = np.zeros_like(w)
    batch_rng = np.random.default_rng(2026)
    # Divide by training input energy: avoids changing the loss scale merely
    # because the all-zero example becomes more common. It preserves minimizers.
    energy = (train_x*train_x).mean()
    for t in range(1, steps+1):
        x = train_x[batch_rng.integers(0, len(train_x), size=256)]
        h = w @ x.T
        pre = h[:, :, None]*w[:, None, :]
        prediction = np.where(relu[:, None, None], np.maximum(pre, 0), pre)
        d = (prediction-x[None, :, :])/(len(x)*energy)
        d = np.where(relu[:, None, None], d*(pre > 0), d)
        dh = (d*w[:, None, :]).sum(axis=2)
        grad = (d*h[:, :, None] + dh[:, :, None]*x[None, :, :]).sum(axis=1)
        moment = .9*moment+.1*grad
        variance = .999*variance+.001*grad*grad
        lr = .0002+.0198*.5*(1+math.cos(math.pi*(t-1)/steps))
        w -= lr*(moment/(1-.9**t))/(np.sqrt(variance/(1-.999**t))+1e-8)
    return initial, w


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--steps', type=int, default=3000)
    parser.add_argument('--check-only', action='store_true')
    parser.add_argument('--output', type=Path, default=Path(__file__).resolve().parents[1]/'data/results.json')
    args = parser.parse_args()
    checks()
    if args.check_only:
        return
    if args.steps < 1:
        parser.error('--steps must be positive')
    groups = []
    for pi, p in enumerate(PS):
        for mi, mode in enumerate(MODES):
            condition_seed = 1000+pi*10+mi
            train_x, _, probabilities = data(p, mode, 8192, condition_seed)
            test_x, states, _ = data(p, mode, 32768, condition_seed+10000)
            initial, final = train(train_x, args.steps)
            # Exact best rank-1 tied linear reconstruction on the training set.
            moment = train_x.T @ train_x/len(train_x)
            _, vectors = np.linalg.eigh(moment)
            optimal_w = vectors[:, -1]
            baseline = metrics(test_x, states, optimal_w, False)
            runs = []
            for idx, w in enumerate(final):
                relu = idx >= len(SEEDS)
                init_idx = idx % len(SEEDS)
                train_mse = float(((forward(train_x, w, relu)-train_x)**2).mean())
                if not relu:
                    optimal_train_mse = ((forward(train_x, optimal_w, False)-train_x)**2).mean()
                    assert train_mse+1e-10 >= optimal_train_mse
                probes = np.array([[1.,0.],[0.,1.],[1.,1.],[0.,0.]])
                runs.append({'seed':SEEDS[init_idx], 'decoder':'relu' if relu else 'linear',
                    'initial_weights':initial[init_idx].tolist(), 'weights':w.tolist(),
                    'opposite_signs':bool(w[0]*w[1]<0), 'train_mse':train_mse,
                    'probe_outputs':forward(probes, w, relu).tolist(),
                    **metrics(test_x, states, w, relu)})
            summary = {}
            for decoder in ['linear','relu']:
                subset = [run for run in runs if run['decoder']==decoder]
                vals = [run['relative_error'] for run in subset]
                summary[decoder] = {'mean':float(np.mean(vals)), 'min':min(vals), 'max':max(vals),
                    'opposite_sign_runs':sum(run['opposite_signs'] for run in subset)}
            groups.append({'p':p,'mode':mode,'probabilities':probabilities,'data_seed':condition_seed,
                'summary':summary,'linear_optimum':{'weights':optimal_w.tolist(), **baseline},'runs':runs})
            print(f'p={p:g} {mode}: linear={summary["linear"]["mean"]:.4f}; relu={summary["relu"]["mean"]:.4f}; opposite={summary["relu"]["opposite_sign_runs"]}/5', flush=True)
    results={'protocol':{'numpy_version':np.__version__,'steps':args.steps,'seeds':SEEDS,
        'train_examples':8192,'test_examples':32768,'batch_size':256,
        'amplitude_range':[.5,1.5],'optimizer':'Adam, cosine learning rate 0.02 to 0.0002',
        'model':'h = w @ x; prediction = w * h, optionally followed by ReLU; no biases; tied weights',
        'normalization':'Training MSE divided by training input energy; relative test error uses test input energy',
        'selection':'All five initializations retained; no early stopping or test-based run selection'},'groups':groups}
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(results,indent=2,allow_nan=False)+'\n')
    print(f'Wrote {args.output}',flush=True)

if __name__=='__main__':
    main()