"""Run: python pandas_demo.py --output results; requires pandas, matplotlib."""
from pathlib import Path
from io import StringIO
import argparse
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
DATA='''id,team,score
001,A,80
002,A,90
003,B,70
004,B,
005,B,0
'''
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)
    frame=pd.read_csv(StringIO(DATA),dtype={'id':'string'})
    frame['score']=pd.to_numeric(frame['score'],errors='coerce')
    frame['passed']=frame['score'].ge(60).where(frame['score'].notna()).astype('boolean')
    summary=frame.groupby('team')['score'].agg(['size','count','mean']).reset_index()
    assert frame.loc[0,'id']=='001'
    assert summary['count'].tolist()==[2,2] and summary['mean'].tolist()==[85,35]
    labels=pd.DataFrame({'team':['A','B'],'mentor':['Lin','Chen']})
    joined=frame.merge(labels,on='team',how='left',validate='many_to_one')
    assert len(joined)==5 and joined['mentor'].notna().all()
    joined.to_csv(args.output/'students.csv',index=False);summary.to_csv(args.output/'summary.csv',index=False)
    fig,axs=plt.subplots(1,2,figsize=(8,3.5),layout='constrained')
    summary.plot.bar(x='team',y=['size','count'],ax=axs[0],rot=0,color=['#adcfe0','#1688b8']);axs[0].set(title='Rows vs valid scores',ylabel='Count')
    summary.plot.bar(x='team',y='mean',ax=axs[1],rot=0,legend=False,color='#1688b8');axs[1].set(title='Means exclude missing scores',ylabel='Score',ylim=(0,100))
    fig.savefig(args.output/'pandas-summary.png',dpi=160);plt.close(fig)
    print(summary.to_string(index=False))
if __name__=='__main__':main()