"""教学模拟，只有显式运行才写指定成果目录，不读取研究数据、不联网或安装软件。"""
from pathlib import Path
import argparse,csv,json,platform
import numpy as np
SEED=42
def generate_sample(n: int = 30, seed: int = SEED, *, heteroskedastic: bool = False):
    """返回 x/y/真实误差；两种生成方式均保持独立观测、零条件均值和已知斜率 2。"""
    if n < 3:
        raise ValueError("至少需要3个观察，才能估计截距/斜率并计算标准误。")
    rng = np.random.default_rng(seed)
    x = rng.uniform(-2, 2, n) if heteroskedastic else rng.normal(0, 1, n)
    sigma = 0.3 + 1.2 * np.abs(x) if heteroskedastic else np.ones(n)
    error = sigma * rng.normal(0, 1, n)
    return x, 1 + 2 * x + error, error

def _bayes_inputs(x, y, prior_mean, prior_sd, sigma, intercept):
    """校验一维教学模型；允许任意有限 x，正规先验保证后验存在。"""
    x, y = np.asarray(x, dtype=float), np.asarray(y, dtype=float)
    if x.ndim != 1 or y.shape != x.shape or x.size == 0:
        raise ValueError("x/y 须为非空等长一维数组。")
    if not np.isfinite(x).all() or not np.isfinite(y).all():
        raise ValueError("教学模型不接受缺失或无穷值。")
    if not np.isfinite([prior_mean, prior_sd, sigma, intercept]).all() or prior_sd <= 0 or sigma <= 0:
        raise ValueError("先验标准差和已知误差标准差须为有限正数，均值和截距须有限。")
    return x, y

def beta_posterior(x, y, *, prior_mean=0.0, prior_sd=5.0, sigma=1.0, intercept=1.0):
    """固定截距/误差方差的正态回归斜率后验；解析结果仅作抽样核对参照。"""
    x, y = _bayes_inputs(x, y, prior_mean, prior_sd, sigma, intercept)
    variance = 1 / (1 / prior_sd**2 + (x @ x) / sigma**2)
    mean = variance * (prior_mean / prior_sd**2 + x @ (y - intercept) / sigma**2)
    return {"mean": float(mean), "sd": float(np.sqrt(variance))}

def metropolis_beta(x, y, *, step_sd=0.35, iterations=5000, seed=SEED,
                    initial=0.0, prior_mean=0.0, prior_sd=5.0, sigma=1.0, intercept=1.0):
    """用对称正态随机游走采样未归一化后验，返回每次迭代及全程接受率。"""
    x, y = _bayes_inputs(x, y, prior_mean, prior_sd, sigma, intercept)
    if not isinstance(iterations, int) or isinstance(iterations, bool) or iterations < 1:
        raise ValueError("iterations 须为正整数。")
    if not np.isfinite(step_sd) or step_sd <= 0 or not np.isfinite(initial):
        raise ValueError("proposal 标准差须为有限正数，初值须有限。")
    rng = np.random.default_rng(seed)

    def log_target(beta):
        residual = (y - intercept - beta * x) / sigma
        return -0.5 * (residual @ residual) - 0.5 * ((beta - prior_mean) / prior_sd)**2

    current = float(initial)
    current_logp = log_target(current)
    chain = np.empty(iterations)
    accepted = 0
    flags = np.zeros(iterations, dtype=int)
    for i in range(iterations):
        proposal = current + rng.normal(0, step_sd)
        proposal_logp = log_target(proposal)
        # 对称 proposal 的 Hastings 比率为1；用 log 比较避免直接指数的数值溢出。
        if np.log(rng.random()) < min(0.0, proposal_logp - current_logp):
            current, current_logp = float(proposal), proposal_logp
            accepted += 1
            flags[i] = 1
        # 拒绝也必须保存当前状态；只留下接受点会改变目标分布及课堂解释。
        chain[i] = current
    return {"chain": chain, "acceptance_rate": accepted / iterations, "accepted": flags}

def run(output,seed=42,iterations=5000,burn=1000):
    from statistics import NormalDist
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    if not 0<=burn<iterations:raise ValueError('Require 0 <= burn < iterations')
    output.mkdir(parents=True,exist_ok=True)
    x,y,error=generate_sample(30,seed);target=beta_posterior(x,y);steps=[.01,.35,3.0]
    with (output/'simulated-data.csv').open('w',newline='') as f:
        w=csv.writer(f,lineterminator='\n');w.writerow(['x','y','error']);w.writerows(zip(x,y,error))
    runs=[metropolis_beta(x,y,step_sd=s,iterations=iterations,seed=seed) for s in steps];kept=[r['chain'][burn:] for r in runs]
    with (output/'chain-states.csv').open('w',newline='') as f:
        w=csv.writer(f,lineterminator='\n');w.writerow(['step_sd','iteration','state','accepted','retained'])
        for step,r in zip(steps,runs):w.writerows((step,i+1,float(v),int(r['accepted'][i]),int(i>=burn)) for i,v in enumerate(r['chain']))
    rows=[{'step_sd':s,'iterations':iterations,'discarded':burn,'retained':len(v),'acceptance_rate':r['acceptance_rate'],'mean':float(v.mean()),'sd':float(v.std(ddof=1)) if len(v)>1 else None,'quantile_025':float(np.quantile(v,.025)),'quantile_975':float(np.quantile(v,.975))} for s,r,v in zip(steps,runs,kept)]
    with (output/'summary.csv').open('w',newline='') as f:
        w=csv.DictWriter(f,fieldnames=list(rows[0]),lineterminator='\n');w.writeheader();w.writerows(rows)
    grid=np.linspace(min(min(v) for v in kept)-.1,max(max(v) for v in kept)+.1,500);density=np.exp(-.5*((grid-target['mean'])/target['sd'])**2)/(target['sd']*np.sqrt(2*np.pi))
    bins=np.linspace(grid.min(),grid.max(),40)
    fig,axes=plt.subplots(3,2,figsize=(12,10),layout='constrained')
    low=min(0,min(min(r['chain']) for r in runs));high=max(max(r['chain']) for r in runs)
    peak=max(max(density),*(max(np.histogram(v,bins=bins,density=True)[0]) for v in kept))
    for row,(s,r,v) in enumerate(zip(steps,runs,kept)):
        axes[row,0].plot(np.arange(1,iterations+1),r['chain'],color='#615b49',linewidth=.6);axes[row,0].axvspan(0,burn,color='#D4C3A1',alpha=.4);axes[row,0].axhline(target['mean'],color='#8c6e36',linestyle='--')
        axes[row,0].set(title=f"Step SD={s:g}; acceptance={r['acceptance_rate']:.1%}",xlabel='Iteration',ylabel='Slope state',ylim=(low-.1,high+.1))
        axes[row,1].hist(v,bins=bins,density=True,color='#D4C3A1',alpha=.8);axes[row,1].plot(grid,density,color='#615b49',label='Analytic posterior');axes[row,1].legend(frameon=False)
        axes[row,1].set(title=f'After discarding {burn} states',xlabel='Slope',ylabel='Density',ylim=(0,peak*1.1))
    for a in axes.flat:a.grid(alpha=.15)
    for ext in ['svg','png']:fig.savefig(output/f'mcmc-proposal-scales.{ext}',dpi=160)
    plt.close(fig)
    svg=output/'mcmc-proposal-scales.svg';svg.write_text('\n'.join(l.rstrip() for l in svg.read_text().splitlines())+'\n')
    q=NormalDist().inv_cdf(.975);target.update(quantile_025=target['mean']-q*target['sd'],quantile_975=target['mean']+q*target['sd'])
    result={'simulated':True,'seed':seed,'n':30,'fixed_intercept':1,'known_sigma':1,'prior_mean':0,'prior_sd':5,'initial':0,'iterations':iterations,'discarded':burn,'analytic':target,'runs':rows,'runtime':{'python':platform.python_version(),'numpy':np.__version__,'matplotlib':matplotlib.__version__}}
    (output/'result.json').write_text(json.dumps(result,indent=2)+'\n');return result
if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('--seed',type=int,default=42);p.add_argument('--iterations',type=int,default=5000);p.add_argument('--discard',type=int,default=1000);p.add_argument('--output-dir',type=Path,default=Path(__file__).resolve().parent/'results');a=p.parse_args()
    print(json.dumps(run(a.output_dir,a.seed,a.iterations,a.discard),indent=2));print('MCMC_SCALES_COMPLETED')