"""Run: python numpy_demo.py --output results; requires numpy, matplotlib."""
from pathlib import Path
import argparse
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

def main():
    p=argparse.ArgumentParser();p.add_argument('--output',type=Path,default=Path('results'));args=p.parse_args();args.output.mkdir(parents=True,exist_ok=True)
    a=np.arange(12,dtype=np.float64).reshape(3,4)
    row_mean=a.mean(axis=1,keepdims=True)
    centered=a-row_mean
    assert a.shape==(3,4) and row_mean.shape==(3,1)
    np.testing.assert_allclose(centered.mean(axis=1),0)
    view=a[:,1:3];view[0,0]=99
    assert a[0,1]==99
    a[0,1]=1
    copied=a[:,1:3].copy();copied[0,0]=-1;assert a[0,1]==1
    fig,axes=plt.subplots(1,3,figsize=(10,3.5))
    for ax,title,data in zip(axes,['a: shape (3, 4)','row means: shape (3, 1)','a - row_mean: (3, 4)'],[a,row_mean,centered]):
        ax.imshow(data,cmap='Blues' if title.startswith('a:') else 'coolwarm',vmin=-5,vmax=11)
        for (r,c),value in np.ndenumerate(data):ax.text(c,r,f'{value:g}',ha='center',va='center',color='black')
        ax.set_title(title,fontsize=10);ax.set_xticks(range(data.shape[1]));ax.set_yticks(range(data.shape[0]))
    fig.tight_layout();fig.savefig(args.output/'numpy-broadcast.png',dpi=160);plt.close(fig)
    np.save(args.output/'centered.npy',centered)
    np.testing.assert_allclose(np.load(args.output/'centered.npy'),centered)
    print('shape:',a.shape,'dtype:',a.dtype);print('row means:',row_mean.ravel().tolist());print('centered row means:',centered.mean(axis=1).tolist())
if __name__=='__main__':main()