import json
from pathlib import Path
import numpy as np

# Fixed-length categorical NAS encoding: 8 operation tokens.
# 0=skip, 1=3x3 conv, 2=1x1 conv, 3=pool.
N_TOKENS, N_OPS = 8, 4
TOKEN_PRIOR = np.array([0.23, 0.38, 0.25, 0.14], dtype=float)
BUDGET = 28.0


def params(a):
    return float(np.array([0.0, 5.0, 3.0, 1.5])[np.asarray(a)].sum())


def latency(a):
    return float(np.array([0.25, 1.4, 0.9, 0.8])[np.asarray(a)].sum())


def valid(a):
    a = np.asarray(a)
    return bool(np.count_nonzero(a != 0) >= 2 and
                not np.any((a[:-1] == 3) & (a[1:] == 3)) and
                params(a) <= BUDGET and latency(a) <= 10.0)


def accuracy(a):
    a = np.asarray(a)
    return float(70 + 2.1*np.count_nonzero(a == 1) + np.count_nonzero(a == 2)
                 - .7*np.count_nonzero(a == 3)
                 + .8*np.sin(np.dot(a + 1, np.arange(1, N_TOKENS + 1))))


def sample_prior(rng, n):
    out = []
    while len(out) < n:
        for a in rng.choice(N_OPS, size=(max(64, n), N_TOKENS), p=TOKEN_PRIOR):
            if valid(a):
                out.append(a.copy())
                if len(out) == n: break
    return np.asarray(out, dtype=int)


def repair(a, rng):
    """Toy external graph checker: repair local structural/resource violations."""
    a = np.asarray(a, dtype=int).copy()
    for i in range(N_TOKENS - 1):
        if a[i] == 3 and a[i+1] == 3:
            a[i+1] = int(rng.choice([0, 1, 2], p=[.25, .5, .25]))
    while np.count_nonzero(a != 0) < 2:
        a[int(rng.integers(N_TOKENS))] = int(rng.choice([1, 2], p=[.65, .35]))
    while params(a) > BUDGET:
        inds = np.where(a == 1)[0]
        if len(inds): a[int(inds[-1])] = 2
        else:
            inds = np.where(a == 2)[0]
            if len(inds): a[int(inds[-1])] = 0
            else: break
    if not valid(a):
        a[:] = 0; a[0], a[1] = 1, 1
    return a


def mutate(parent, gamma, rng, do_repair=True):
    """Partial categorical forward noising plus a simple conditional reverse kernel.

    For t=round(gamma*T), each token is retained with alpha_bar_t and otherwise
    sampled from the learned categorical marginal. This is the exact categorical
    analogue of the stated anchored forward step; reverse sampling leaves the
    retained evidence anchored and resamples corrupted positions.
    """
    T = 20; t = int(round(gamma*T))
    betas = np.linspace(.01, .30, T)
    alpha_bar = float(np.prod(1 - betas[:t])) if t else 1.0
    child = np.asarray(parent, dtype=int).copy()
    changed = rng.random(N_TOKENS) > alpha_bar
    child[changed] = rng.choice(N_OPS, size=int(changed.sum()), p=TOKEN_PRIOR)
    raw = child.copy()
    return (repair(raw, rng) if do_repair else raw), raw, alpha_bar


def edit(a, b):
    return float(np.mean(np.asarray(a) != np.asarray(b)))


def main():
    seed = 2761; rng = np.random.default_rng(seed)
    parents = sample_prior(rng, 200)
    gammas = [0., .05, .15, .30, .60, 1.]
    rows = []
    for g in gammas:
        repaired, raw = [], [] ; alphas = []
        for p in parents:
            c, r, ab = mutate(p, g, rng, True)
            repaired.append(c); raw.append(r); alphas.append(ab)
        repaired, raw = np.asarray(repaired), np.asarray(raw)
        raw_valid = np.mean([valid(x) for x in raw])
        rows.append({
            'gamma': g, 'alpha_bar': float(np.mean(alphas)),
            'mean_edit': float(np.mean([edit(p,c) for p,c in zip(parents,repaired)])),
            'raw_validity': float(raw_valid), 'post_repair_validity': float(np.mean([valid(x) for x in repaired])),
            'mean_accuracy': float(np.mean([accuracy(x) for x in repaired])),
            'p95_edit': float(np.quantile([edit(p,c) for p,c in zip(parents,repaired)], .95))})

    indep = sample_prior(rng, len(parents))
    indep_edit = float(np.mean([edit(p,c) for p,c in zip(parents,indep)]))
    result = {
      'seed': seed, 'rows': rows,
      'independent': {'mean_edit': indep_edit, 'validity': 1.0},
      'retention_check': [],
      'predictions': {
        'P1': 'categorical q retains each token with alpha_bar; observed retention should match alpha_bar within 0.01',
        'P2': 'edit distance should increase monotonically with gamma',
        'P3': 'gamma=1 mutation edit distance should approach independent-prior edit distance; gap expected small relative to 1 token fraction'}}
    crng = np.random.default_rng(991)
    for g in gammas[1:]:
        t = int(round(g*20)); ab = float(np.prod(1-np.linspace(.01,.30,20)[:t]))
        observed = float((crng.random((5000,N_TOKENS)) < ab).mean())
        result['retention_check'].append({'gamma':g, 'predicted_retention':ab, 'observed_retention':observed,
                                          'absolute_error':abs(ab-observed)})
    result['convergence'] = {'gamma1_vs_independent_absolute_edit_gap': abs(rows[-1]['mean_edit']-indep_edit),
                             'gamma005_vs_independent_absolute_edit_gap': abs(rows[1]['mean_edit']-indep_edit)}
    Path('results.json').write_text(json.dumps(result, indent=2)); print(json.dumps(result, indent=2))

if __name__ == '__main__': main()