#!/usr/bin/env python3
"""Package each retained AI stack candidate as an explicit 49-member SPEI product."""
import argparse,json
from datetime import datetime
from pathlib import Path
import numpy as np
import spei_pipeline as core
ROOT=core.ROOT; METHODS=core.METHODS

def main():
 p=argparse.ArgumentParser();p.add_argument('--run',required=True);a=p.parse_args();base=ROOT/'outputs/production_output/spei_members'/a.run
 lon,lat,_=core.coordinates_and_mask();manifest=[]
 for lead in range(1,7):
  for scale in (1,3,6):
   members=core.nc_array(base/f'SPEI{scale}_members_LT{lead}.nc',f'SPEI{scale}')
   stacks=core.nc_array(base/f'SPEI{scale}_AI_stack_candidates_LT{lead}.nc',f'SPEI{scale}_stack')
   assert members.shape==(40,68,48,480),members.shape;assert stacks.shape==(40,68,4,480),stacks.shape
   for mi,method in enumerate(METHODS):
    data=np.concatenate([members,stacks[:,:,mi:mi+1,:]],axis=2);assert data.shape==(40,68,49,480)
    names=[f'SPEI_member_{i:02d}' for i in range(1,49)]+[f'AI_Stack_{method}']
    path=base/f'SPEI{scale}_49members_{method}_LT{lead}.nc'
    core.write_nc(path,f'SPEI{scale}',data,lon,lat,names,{'source_step_1_5_run':'1399282','period':'1982-01/2021-12','member_count':49,'members_1_48':'SPEI from 48 step-1.5 pr/t2 member pairs','member_49':f'AI Stack {method}','best_model_selected':'false','lead':lead,'timescale_months':scale})
    manifest.append({'file':path.name,'lead':lead,'timescale':scale,'stack_method':method,'shape':list(data.shape),'members':49})
 (base/'49_member_manifest.json').write_text(json.dumps({'run':a.run,'created_at':datetime.now().isoformat(),'source_step_1_5_run':'1399282','best_model_selected':False,'products':manifest},indent=2))
 assert len(manifest)==72
 (base/'SUCCESS_49_MEMBERS').write_text(datetime.now().isoformat())
if __name__=='__main__':main()