"""OLS 教学：只生成模拟数据，输出指定练习目录的 CSV/图/摘要，不读研究文件。
复用官网已公开的样本与重复抽样逻辑；CLI 显式执行才写文件，不联网或安装包。
"""
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 fit_ols(x, y):
    """含截距的单变量OLS，返回系数、普通/HC1标准误、拟合值和残差。"""
    x, y = np.asarray(x, dtype=float), np.asarray(y, dtype=float)
    if x.ndim != 1 or y.shape != x.shape or len(x) <= 2:
        raise ValueError("x/y须为等长一维数组，且观察数大于2。")
    if not np.isfinite(x).all() or not np.isfinite(y).all():
        raise ValueError("示例不接受缺失或无穷值；先明确样本处理规则。")
    design = np.column_stack([np.ones(len(x)), x])
    beta, _, rank, _ = np.linalg.lstsq(design, y, rcond=None)
    if rank != 2:
        raise ValueError("解释变量必须有变化，否则无法识别斜率。")
    fitted = design @ beta
    residual = y - fitted
    n, k = design.shape
    bread = np.linalg.inv(design.T @ design)
    classical_cov = (residual @ residual) / (n - k) * bread
    # HC1只调整协方差估计，不重估系数；n/(n-k)是有限样本修正，k包括截距。
    scores = design * residual[:, None]
    hc1_cov = n / (n - k) * bread @ (scores.T @ scores) @ bread
    return {
        "beta": beta,
        "classical_se": np.sqrt(np.diag(classical_cov)),
        "hc1_se": np.sqrt(np.diag(hc1_cov)),
        "fitted": fitted,
        "residual": residual,
    }


def sampling_experiment(repetitions: int = 500, seed: int = SEED):
    """每次重新抽取独立样本；同一轮的30与300并非一个样本的前缀。"""
    if repetitions < 2:
        raise ValueError("重复抽样至少需要2次。")
    rng = np.random.default_rng(seed)
    slopes = {}
    for n in (30, 300):
        estimates = np.empty(repetitions)
        for i in range(repetitions):
            x = rng.normal(0, 1, n)
            y = 1 + 2 * x + rng.normal(0, 1, n)
            design = np.column_stack([np.ones(n), x])
            estimates[i] = np.linalg.lstsq(design, y, rcond=None)[0][1]
        slopes[n] = estimates
    return slopes



def run(output, seed=SEED, repetitions=500):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    output.mkdir(parents=True,exist_ok=True)
    x,y,error=generate_sample(30,seed)
    fit=fit_ols(x,y)
    slopes=sampling_experiment(repetitions,seed)
    with (output/'single-sample.csv').open('w',newline='') as stream:
        writer=csv.writer(stream,lineterminator='\n');writer.writerow(['x','y','error','fitted','residual'])
        writer.writerows(zip(x,y,error,fit['fitted'],fit['residual']))
    with (output/'sampling-slopes.csv').open('w',newline='') as stream:
        writer=csv.writer(stream,lineterminator='\n');writer.writerow(['n','repetition','slope'])
        for n,values in slopes.items():
            writer.writerows((n,i+1,float(v)) for i,v in enumerate(values))
    fig,axes=plt.subplots(1,2,figsize=(12,5),layout='constrained')
    axes[0].scatter(x,y,color='#615b49',s=28)
    grid=np.linspace(x.min(),x.max(),120)
    axes[0].plot(grid,1+2*grid,color='#333333',linestyle='--',label='True slope = 2')
    axes[0].plot(grid,fit['beta'][0]+fit['beta'][1]*grid,color='#8c6e36',label=f"OLS slope = {fit['beta'][1]:.3f}")
    axes[0].set(title='One simulated sample (n=30)',xlabel='x (simulated units)',ylabel='y (simulated units)')
    # 同一 bins/密度单位便于比较，不将两种样本量当作同一批观察的前缀。
    bins=np.linspace(min(v.min() for v in slopes.values()),max(v.max() for v in slopes.values()),36)
    axes[1].hist(slopes[30],bins=bins,density=True,color='#D4C3A1',alpha=.8,label='n=30')
    axes[1].hist(slopes[300],bins=bins,density=True,histtype='step',linewidth=2,color='#615b49',label='n=300')
    axes[1].axvline(2,color='#333333',linestyle='--',label='True slope = 2')
    axes[1].set(title=f'{repetitions} independent samples per size',xlabel='Estimated slope',ylabel='Density')
    for ax in axes:
        ax.legend(frameon=False);ax.grid(alpha=.15)
    for ext in ['png','svg']: fig.savefig(output/f'ols-sampling.{ext}',dpi=160)
    plt.close(fig)
    svg=output/'ols-sampling.svg';svg.write_text('\n'.join(line.rstrip() for line in svg.read_text().splitlines())+'\n')
    summary={'simulated':True,'seed':seed,'true_intercept':1,'true_slope':2,'single_n':30,'single_intercept':float(fit['beta'][0]),'single_slope':float(fit['beta'][1]),'sampling':[{'n':n,'repetitions':len(v),'mean_slope':float(v.mean()),'sd_slope':float(v.std(ddof=1)),'quantile_025':float(np.quantile(v,.025)),'quantile_975':float(np.quantile(v,.975))} for n,v in slopes.items()],'runtime':{'python':platform.python_version(),'numpy':np.__version__,'matplotlib':matplotlib.__version__}}
    (output/'result.json').write_text(json.dumps(summary,indent=2)+'\n')
    return summary

if __name__=='__main__':
    parser=argparse.ArgumentParser(description='Generate an entirely simulated OLS teaching exercise.')
    parser.add_argument('--output-dir',type=Path,default=Path(__file__).resolve().parent/'results')
    parser.add_argument('--seed',type=int,default=42)
    parser.add_argument('--repetitions',type=int,default=500)
    args=parser.parse_args()
    print(json.dumps(run(args.output_dir,args.seed,args.repetitions),indent=2))
    print('OLS_SAMPLING_COMPLETED')