import random, json, platform
from collections import Counter
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

# Change one variable: threshold. At <=0.5 uncertain positions commit together.
COMMIT_THRESHOLD=.99
VISUAL_SEED=20260928
VISUAL_N=4000
OUT=Path('lab04-visual-output');OUT.mkdir(exist_ok=True)
GRAMMAR_VIS={('cat','can','meow',',','cat','meow'):.5,('dog','can','woof',',','dog','woof'):.5}

def vis_conditionals(state):
    compatible={s:p for s,p in GRAMMAR_VIS.items() if all(v is None or v==s[i] for i,v in enumerate(state))}
    total=sum(compatible.values())
    if not total:raise ValueError('Visible tokens conflict with the fixed grammar')
    out={}
    for i,v in enumerate(state):
        if v is None:
            d=Counter()
            for s,p in compatible.items():d[s[i]]+=p/total
            out[i]=dict(d)
    return out

def vis_draw(d,rng):return rng.choices(list(d),weights=list(d.values()),k=1)[0]
def vis_sampler(mode,seed,threshold=COMMIT_THRESHOLD,initial=None):
    rng=random.Random(seed);state=list(initial) if initial else [None]*6;trace=[state.copy()];new=[]
    while None in state:
        dist=vis_conditionals(state)
        if mode=='parallel':selected=list(dist)
        elif mode=='AR':selected=[min(dist)]
        else:
            selected=[i for i,p in dist.items() if max(p.values())>=threshold]
            if not selected:selected=[rng.choice(list(dist))]
        updates={i:vis_draw(dist[i],rng) for i in selected}
        for i,value in updates.items():state[i]=value
        trace.append(state.copy());new.append(selected)
    return {'tokens':state,'rounds':len(new),'valid':tuple(state) in GRAMMAR_VIS,'trace':trace,'new':new}

fig,axes=plt.subplots(1,3,figsize=(13,5.2));traces={}
for ax,mode in zip(axes,['parallel','iterative','AR']):
    r=vis_sampler(mode,2);traces[mode]=r
    for row,state in enumerate(r['trace']):
        changed=[] if row==0 else r['new'][row-1]
        for col,value in enumerate(state):
            color='#edf0f2' if value is None else '#f8dc8b' if col in changed else '#cbdde8'
            ax.add_patch(Rectangle((col,-row),.94,.8,facecolor=color,edgecolor='white'))
            ax.text(col+.47,-row+.4,'?' if value is None else value,ha='center',va='center',fontsize=9)
        ax.text(-.15,-row+.4,str(row),ha='right',va='center',fontsize=9)
    ax.set(xlim=(-.5,6),ylim=(-6.5,1),title=f'{mode}: {r["rounds"]} rounds; valid={r["valid"]}');ax.axis('off')
fig.suptitle('Same fully masked input: gray unknown, gold newly committed, blue retained',fontsize=12)
fig.tight_layout();fig.savefig(OUT/'lab04-token-traces.png',dpi=160);plt.show();plt.close(fig)

records=[]
for mode in ['parallel','iterative','AR']:
    outputs=[vis_sampler(mode,VISUAL_SEED+i) for i in range(VISUAL_N)]
    records.append({'mode':mode,'valid_rate':sum(r['valid'] for r in outputs)/VISUAL_N,'mean_rounds':float(np.mean([r['rounds'] for r in outputs])),'cat_mass':sum(tuple(r['tokens'])==next(iter(GRAMMAR_VIS)) for r in outputs)/VISUAL_N})
fig,axes=plt.subplots(1,2,figsize=(10,3.6))
axes[0].bar([r['mode'] for r in records],[r['valid_rate'] for r in records],color=['#a3b4c2','#003262','#c18f16']);axes[0].axhline(.125,color='#a5563c',ls='--',label='one-shot exact 1/8');axes[0].set(ylim=(0,1.08),ylabel='Joint validity');axes[0].legend(fontsize=8)
axes[1].bar([r['mode'] for r in records],[r['mean_rounds'] for r in records],color=['#a3b4c2','#003262','#c18f16']);axes[1].set(ylabel='Dependency rounds (not wall-clock time)',ylim=(0,7))
fig.suptitle(f'{VISUAL_N} seeds; iterative threshold={COMMIT_THRESHOLD:g}',fontsize=12);fig.tight_layout();fig.savefig(OUT/'lab04-validity-rounds.png',dpi=160);plt.show();plt.close(fig)
assert abs(records[0]['valid_rate']-.125)<.03
if COMMIT_THRESHOLD>.5:assert records[1]['valid_rate']==1
assert records[2]['valid_rate']==1
payload={'kind':'exact two-sentence grammar, no language-model inference','seed':VISUAL_SEED,'n':VISUAL_N,'threshold':COMMIT_THRESHOLD,'traces':traces,'results':records,'legend':{'gray':'unknown','gold':'newly committed this round','blue':'retained; no remasking/revision in this scheduler'},'python':platform.python_version()}
(OUT/'lab04-results.json').write_text(json.dumps(payload,indent=2));print(json.dumps(records,indent=2))

[
  {
    "mode": "parallel",
    "valid_rate": 0.1325,
    "mean_rounds": 1.0,
    "cat_mass": 0.06825
  },
  {
    "mode": "iterative",
    "valid_rate": 1.0,
    "mean_rounds": 3.0,
    "cat_mass": 0.49375
  },
  {
    "mode": "AR",
    "valid_rate": 1.0,
    "mean_rounds": 6.0,
    "cat_mass": 0.51
  }
]

import itertools
import random
from collections import Counter

SENTENCES = {
    ("猫", "会", "喵", "，", "猫", "喵"): 0.5,
    ("狗", "会", "汪", "，", "狗", "汪"): 0.5,
}
LENGTH = 6
def show(state):
    return " ".join(token if token is not None else "[MASK]" for token in state)

def conditionals(state):
    compatible = {seq: p for seq,p in SENTENCES.items()
                  if all(token is None or token == seq[i] for i,token in enumerate(state))}
    total = sum(compatible.values())
    if total == 0:
        raise ValueError("Visible tokens contradict every sentence in this fixed model.")
    output = {}
    for i, token in enumerate(state):
        if token is None:
            dist = Counter()
            for seq,p in compatible.items():
                dist[seq[i]] += p/total
            output[i] = dict(dist)
    return output

source = next(iter(SENTENCES))
rng_corrupt = random.Random(10)
mask_rate = 0.7
corrupted = [None if rng_corrupt.random() < mask_rate else t for t in source]
print("clean:", show(source))
print("masked:", show(corrupted))
print("conditional predictions:", conditionals(corrupted))
print("fully masked predictions:", conditionals([None]*LENGTH))

clean: 猫 会 喵 ， 猫 喵
masked: [MASK] [MASK] [MASK] [MASK] 猫 喵
conditional predictions: {0: {'猫': 1.0}, 1: {'会': 1.0}, 2: {'喵': 1.0}, 3: {'，': 1.0}}
fully masked predictions: {0: {'猫': 0.5, '狗': 0.5}, 1: {'会': 1.0}, 2: {'喵': 0.5, '汪': 0.5}, 3: {'，': 1.0}, 4: {'猫': 0.5, '狗': 0.5}, 5: {'喵': 0.5, '汪': 0.5}}

marginals = conditionals([None]*LENGTH)
parallel_joint = {}
for candidate in itertools.product(*(list(marginals[i]) for i in range(LENGTH))):
    probability = 1.0
    for i, token in enumerate(candidate):
        probability *= marginals[i][token]
    parallel_joint[candidate] = probability
valid_mass = sum(p for seq,p in parallel_joint.items() if seq in SENTENCES)
print("parallel support size:", len(parallel_joint))
print("exact valid probability:", valid_mass)
print("one incompatible candidate:", next(show(s) for s in parallel_joint if s not in SENTENCES))
assert abs(sum(parallel_joint.values()) - 1) < 1e-12
assert valid_mass == 0.125

parallel support size: 16
exact valid probability: 0.125
one incompatible candidate: 猫 会 喵 ， 猫 汪

def draw(dist, rng):
    tokens = list(dist)
    return rng.choices(tokens, weights=[dist[t] for t in tokens], k=1)[0]

def parallel_once(rng, initial=None):
    state = list(initial) if initial is not None else [None]*LENGTH
    prediction = conditionals(state)
    for i, dist in prediction.items():
        state[i] = draw(dist, rng)
    return tuple(state), 1, [show(state)]

def iterative(rng, initial=None, threshold=0.99):
    state = list(initial) if initial is not None else [None]*LENGTH
    trace = [show(state)]
    rounds = 0
    while None in state:
        prediction = conditionals(state)
        confident = [i for i,dist in prediction.items() if max(dist.values()) >= threshold]
        selected = confident if confident else [rng.choice(list(prediction))]
        # 所有 selected 都使用本轮旧状态；不会在同一轮内偷偷传递刚生成的 token。
        updates = {i: draw(prediction[i], rng) for i in selected}
        for i, token in updates.items():
            state[i] = token
        rounds += 1
        trace.append(show(state))
    return tuple(state), rounds, trace

def autoregressive(rng, initial=None):
    state = list(initial) if initial is not None else [None]*LENGTH
    trace = [show(state)]
    rounds = 0
    for i in range(LENGTH):
        if state[i] is None:
            state[i] = draw(conditionals(state)[i], rng)
            rounds += 1
            trace.append(show(state))
    return tuple(state), rounds, trace

for name, sampler in [("parallel once", parallel_once), ("iterative", iterative), ("AR", autoregressive)]:
    result, rounds, trace = sampler(random.Random(4))
    print("\n", name, "rounds=", rounds, "valid=", result in SENTENCES)
    for step, state in enumerate(trace):
        print(step, state)

 parallel once rounds= 1 valid= True
0 猫 会 喵 ， 猫 喵

 iterative rounds= 3 valid= True
0 [MASK] [MASK] [MASK] [MASK] [MASK] [MASK]
1 [MASK] 会 [MASK] ， [MASK] [MASK]
2 [MASK] 会 [MASK] ， [MASK] 喵
3 猫 会 喵 ， 猫 喵

 AR rounds= 6 valid= True
0 [MASK] [MASK] [MASK] [MASK] [MASK] [MASK]
1 猫 [MASK] [MASK] [MASK] [MASK] [MASK]
2 猫 会 [MASK] [MASK] [MASK] [MASK]
3 猫 会 喵 [MASK] [MASK] [MASK]
4 猫 会 喵 ， [MASK] [MASK]
5 猫 会 喵 ， 猫 [MASK]
6 猫 会 喵 ， 猫 喵

SAMPLE_COUNT = 4000
print("sampler | valid rate | cat mass | dog mass | mean rounds")
for index, (name, sampler) in enumerate([
    ("parallel once", parallel_once), ("iterative", iterative), ("AR", autoregressive)
]):
    rng = random.Random(20260926 + index)
    counts = Counter()
    round_sum = 0
    for _ in range(SAMPLE_COUNT):
        output, rounds, _ = sampler(rng)
        counts[output] += 1
        round_sum += rounds
    valid = sum(counts[s] for s in SENTENCES) / SAMPLE_COUNT
    cat_mass, dog_mass = [counts[s]/SAMPLE_COUNT for s in SENTENCES]
    print(name, f"{valid:.4f}", f"{cat_mass:.4f}", f"{dog_mass:.4f}", f"{round_sum/SAMPLE_COUNT:.2f}")
    if name == "parallel once":
        assert abs(valid - valid_mass) < 0.025
    else:
        assert valid == 1.0 and abs(cat_mass - 0.5) < 0.04

prompt = ["狗", None, None, None, None, None]
output, rounds, trace = iterative(random.Random(8), initial=prompt)
assert output[0] == "狗" and output in SENTENCES
print("\nFixed condition:", show(prompt), "->", show(output), "rounds=", rounds)
contradictory = ["猫", None, "汪", None, None, None]
try:
    conditionals(contradictory)
except ValueError as error:
    print("Contradiction detected:", str(error))
else:
    raise AssertionError("Contradictory context must not be silently repaired.")

sampler | valid rate | cat mass | dog mass | mean rounds
parallel once 0.1235 0.0625 0.0610 1.00
iterative 1.0000 0.4925 0.5075 3.00
AR 1.0000 0.4943 0.5058 6.00

Fixed condition: 狗 [MASK] [MASK] [MASK] [MASK] [MASK] -> 狗 会 汪 ， 狗 汪 rounds= 1
Contradiction detected: Visible tokens contradict every sentence in this fixed model.

rng_threshold = random.Random(20260930)
counts_threshold = Counter()
for _ in range(SAMPLE_COUNT):
    output, rounds, _ = iterative(rng_threshold, threshold=0.5)
    counts_threshold[output] += 1
    assert rounds == 1
validity_low_threshold = sum(counts_threshold[s] for s in SENTENCES)/SAMPLE_COUNT
print("threshold=0.5: observed valid rate =", round(validity_low_threshold, 4))
print("exact one-shot independent valid rate =", valid_mass)
assert abs(validity_low_threshold-valid_mass) < 0.025
print("All exact support, seeded sampling, condition and scheduler checks passed.")

threshold=0.5: observed valid rate = 0.1185
exact one-shot independent valid rate = 0.125
All exact support, seeded sampling, condition and scheduler checks passed.