"""Reproduce the bounded historical American Gut analysis.
Run from the extracted bundle root: python analysis.py
Inputs are kept byte-for-byte in ag_analysis. Outputs default to ag_results.
Set AG_OUTPUT_DIR to a new relative directory for subsequent runs.
"""
import numpy as np, pandas as pd, scipy, matplotlib
import scipy.spatial.distance as ssd
import scipy.linalg as la
import pathlib, json, hashlib, urllib.request, sys, importlib.metadata, zipfile
EXPECTED_SHA256={
'AG_100nt.biom':'35e474a374acb6bbf53ba7c2d1ba750b5e3c579eed0b8fc325e63f309cb503bb',
'AG_100nt.txt':'e37d1c5b2e17f62c194da231e499611a391bea76e4d15fe2d1a5b55166fcea4b',
'README.md':'a01016dba296950fb798c0dcb577676d9b776bbfa4d6ba06567cbc9030c0d30f',
'metadata_charts.json':'e13dd86e7c20635eafe12c84ea3f582067614a96b71c8c75ffe34968d440bc01'}

# %% Cell 1
BASE='https://raw.githubusercontent.com/biocore/American-Gut/68fd6d4b2fa6aeb5b4f5272c6f1006defe5b160e/'
root=pathlib.Path('ag_analysis'); root.mkdir(exist_ok=True)
manifest=[]
for remote in ['data/AG/AG_100nt.biom','data/AG/AG_100nt.txt','README.md']:
    target=root/pathlib.Path(remote).name
    if not target.exists(): urllib.request.urlretrieve(BASE+remote,target)
    raw=target.read_bytes()
    assert hashlib.sha256(raw).hexdigest()==EXPECTED_SHA256[target.name], 'Input checksum mismatch'
    manifest.append({'url':BASE+remote,'file':target.name,'bytes':len(raw),'sha256':hashlib.sha256(raw).hexdigest()})
meta=pd.read_csv(root/'AG_100nt.txt',sep='\t',dtype=str,keep_default_na=False).set_index('#SampleID')
print('metadata shape',meta.shape)
print('BIOM signature', (root/'AG_100nt.biom').read_bytes()[:60])
for col in ['TYPES_OF_PLANTS','PLANT_PER','COUNTRY','BODY_SITE','AGE_UNIT','SEX','ANTIBIOTIC_SELECT']:
 print(col,meta[col].value_counts().head(18).to_dict())
print('subject missing',meta.HOST_SUBJECT_ID.value_counts().head(4).to_dict())

# %% Cell 2
import h5py
from scipy.sparse import csr_matrix
with h5py.File(root/'AG_100nt.biom','r') as f:
 print('BIOM',dict(f.attrs))
 samples=f['sample/ids'].asstr()[:]
 features=f['observation/ids'].asstr()[:]
 g=f['sample/matrix']; counts=csr_matrix((g['data'][:],g['indices'][:],g['indptr'][:]),shape=(len(samples),len(features)))
 print('taxonomy keys',list(f['observation/metadata']))
print('counts',counts.shape,'overlap',len(set(samples)&set(meta.index)))
print(meta[['HOST_SUBJECT_ID','AGE','STATE']].head(8).to_dict('records'))
print('subject identifiers',meta.HOST_SUBJECT_ID.nunique(),meta.HOST_SUBJECT_ID.str.contains('unknown',case=False).sum())

# %% Cell 3
remote='latex/metadata_charts.json'
target=root/'metadata_charts.json'
if not target.exists(): urllib.request.urlretrieve(BASE+remote,target)
raw=target.read_bytes(); assert hashlib.sha256(raw).hexdigest()==EXPECTED_SHA256[target.name]; manifest.append({'url':BASE+remote,'file':target.name,'bytes':len(raw),'sha256':hashlib.sha256(raw).hexdigest()})
charts=json.loads(raw)
print(str(charts)[:4500])

# %% Cell 4
SEED=20261008
LEVELS=['Less than 5','6 to 10','11 to 20','21 to 30','More than 30']
missing={'','no_data','unknown','not provided','unspecified','nan','none'}
m=meta.loc[samples].copy(); m['sample_id']=m.index
m['depth']=np.asarray(counts.sum(axis=1)).ravel(); m['age']=pd.to_numeric(m.AGE,errors='coerce'); m['bmi']=pd.to_numeric(m.BMI,errors='coerce')
m['country']=m.COUNTRY.str.replace('GAZ:','',regex=False); m.loc[m.country.str.lower().isin(missing),'country']='Missing'
m['plant']=m.TYPES_OF_PLANTS.where(m.TYPES_OF_PLANTS.isin(LEVELS))
m['state']=m.STATE.where(~m.STATE.str.lower().isin(missing))
m['sex']=m.SEX.where(m.SEX.isin(['female','male','other']))
m['antibiotic']=m.ANTIBIOTIC_SELECT.where(m.ANTIBIOTIC_SELECT.isin(['Not in the last year','In the past year','In the past 6 months','In the past month','In the past week']))
m.loc[~np.isfinite(m.bmi)|(m.bmi<=0),'bmi']=np.nan
flow=[('matched',len(m))]
mask=m.BODY_SITE.eq('UBERON:feces'); flow.append(('stool',int(mask.sum())))
mask &= m.age.ge(18)&m.AGE_UNIT.eq('years'); flow.append(('adult_years',int(mask.sum())))
mask &= m.HOST_SUBJECT_ID.str.fullmatch(r'(?:\d+:)?\d+'); flow.append(('valid_subject',int(mask.sum())))
assert np.isfinite(counts.data).all() and (counts.data>=0).all()
mask &= m.depth.ge(1000); flow.append(('depth_at_least_1000',int(mask.sum())))
m['participant']=m.HOST_SUBJECT_ID.str.split(':').str[-1]
elig=m.loc[mask].sort_values(['depth','sample_id'],ascending=[False,True]).drop_duplicates('participant').sort_index()
flow.append(('one_per_participant',len(elig)))
sizes=elig.groupby('country').size().sort_index(); cap=min(600,len(elig)); quota=sizes/sizes.sum()*cap; alloc=np.floor(quota).astype(int)
for key in (quota-alloc).sort_values(ascending=False,kind='stable').index[:cap-int(alloc.sum())]: alloc[key]+=1
rng=np.random.default_rng(SEED)
chosen=[]
for country,k in alloc.items():
 ids=elig.index[elig.country.eq(country)].to_numpy(); chosen.extend(rng.choice(ids,k,replace=False))
cohort=elig.loc[sorted(chosen)].copy(); flow.append(('capped',len(cohort)))
ix=pd.Index(samples).get_indexer(cohort.index); selected_counts=counts[ix]; feat_mask=np.asarray(selected_counts.sum(axis=0)).ravel()>0
selected_counts=selected_counts[:,feat_mask]; rel=selected_counts.toarray().astype(float); rel/=rel.sum(axis=1,keepdims=True)
cohort['row']=np.arange(len(cohort))
ct=pd.crosstab(cohort.country,cohort.plant).reindex(columns=LEVELS,fill_value=0)
print('FLOW',flow,'FEATURES',int(feat_mask.sum()))
print('SUPPORT',ct.to_dict('index'))
print('missing',cohort[['plant','country','state','age','sex','bmi','antibiotic']].isna().sum().to_dict())

# %% Cell 5
from threadpoolctl import threadpool_limits
threadpool_limits(limits=1)
D=ssd.squareform(ssd.pdist(rel,metric='braycurtis'))
assert np.isfinite(D).all() and np.allclose(D,D.T) and np.allclose(np.diag(D),0)
def centered_gram(d):
 a=-0.5*d*d
 return a-a.mean(axis=0)[None,:]-a.mean(axis=1)[:,None]+a.mean()
B=centered_gram(D); original_eigen=la.eigvalsh(B); lingoes=max(0.,-float(original_eigen[0]))
Dc=np.sqrt(D*D+2*lingoes); np.fill_diagonal(Dc,0)
ev,U=la.eigh(centered_gram(Dc)); tol=1e-10*max(1.,ev[-1]); positive=ev>tol
Y=U[:,positive]*np.sqrt(ev[positive]); Y=Y[:,::-1]
geometry={'n':len(Y),'features':rel.shape[1],'negative_eigenvalues':int((original_eigen < -tol).sum()),'minimum_original_eigenvalue':float(original_eigen[0]),'lingoes_c':lingoes,'dimensions':Y.shape[1],'minimum_corrected_eigenvalue':float(ev[0]),'reconstruction_max_error':float(np.max(np.abs(ssd.squareform(ssd.pdist(Y))-Dc))),'axis1_fraction':float(ev[-1]/ev[positive].sum()),'axis2_fraction':float(ev[-2]/ev[positive].sum())}
assert geometry['reconstruction_max_error']<1e-7
def design(df,cols):
 parts=[np.ones((len(df),1))]
 for col in cols:
  if col in ['age','bmi']:
   a=df[col].to_numpy(float); parts.append(((a-a.mean())/a.std())[:,None])
  else: parts.append(pd.get_dummies(df[col].astype(str),drop_first=True,dtype=float).to_numpy())
 return np.column_stack(parts)
def ortho(x):
 assert np.linalg.matrix_rank(x)==x.shape[1], 'rank-deficient design'
 return la.qr(x,mode='economic')[0]
def fl_test(y,x0,x1,strata=None,nperm=999,seed=SEED):
 q0,q1=ortho(x0),ortho(x1); n=len(y); df=x1.shape[1]-x0.shape[1]; rdf=n-x1.shape[1]
 assert df>0 and rdf>=20
 fit=q0@(q0.T@y); residual=y-fit; ss0=float(np.sum(residual**2))
 ss1=float(np.sum(y*y)-np.sum((q1.T@y)**2)); extra=ss0-ss1
 obs=(extra/df)/(ss1/rdf); total=float(np.sum((y-y.mean(axis=0))**2))
 rng=np.random.default_rng(seed); blocks=[np.arange(n)] if strata is None else [np.flatnonzero(np.asarray(strata)==g) for g in sorted(set(strata))]
 ge=0
 for _ in range(nperm):
  perm=np.arange(n)
  for ids in blocks: perm[ids]=rng.permutation(ids)
  yp=fit+residual[perm]
  norm=float(np.sum(yp*yp)); s0=norm-float(np.sum((q0.T@yp)**2)); s1=norm-float(np.sum((q1.T@yp)**2))
  fp=((s0-s1)/df)/(s1/rdf)
  ge+=fp>=obs-1e-12
 return {'n':n,'df':df,'residual_df':rdf,'pseudo_F':obs,'incremental_R2':extra/total,'partial_R2':extra/ss0,'p':(ge+1)/(nperm+1),'permutations':nperm}
# Numerical identities against ordinary univariate regression; deterministic null permutations.
rng=np.random.default_rng(91); ty=rng.normal(size=(60,1)); tx0=np.ones((60,1)); tx1=np.column_stack([tx0,np.tile([0,1],30)])
test=fl_test(ty,tx0,tx1,nperm=99)
b=np.linalg.lstsq(tx1,ty,rcond=None)[0]; rss=float(np.sum((ty-tx1@b)**2)); tss=float(np.sum((ty-ty.mean())**2))
assert np.isclose(test['pseudo_F'],((tss-rss)/1)/(rss/58))
assert test==fl_test(ty,tx0,tx1,nperm=99)
assert np.isclose(test['incremental_R2'],test['partial_R2'])
main=cohort.loc[cohort.plant.notna()&cohort.country.ne('Missing')].copy()
results=[]
def model(label,df,nuisance,term='plant',stratum=None,family='context'):
 yy=Y[df.row.to_numpy()]
 row=fl_test(yy,design(df,nuisance),design(df,nuisance+[term]),None if stratum is None else df[stratum].to_numpy())
 row.update(model=label,family=family); results.append(row); return row
model('Main unadjusted',main,[])
model('Main country-adjusted',main,['country'],stratum='country',family='primary')
print('geometry',geometry)
print(pd.DataFrame(results).to_string(index=False))


# %% Cell 6
def dispersion(df):
 yy=Y[df.row.to_numpy()]; groups=df.plant.to_numpy(); z=np.empty(len(df))
 for g in sorted(set(groups)):
  ids=np.flatnonzero(groups==g); z[ids]=np.linalg.norm(yy[ids]-yy[ids].mean(axis=0),axis=1)*np.sqrt(len(ids)/(len(ids)-1))
 out=fl_test(z[:,None],np.ones((len(df),1)),design(df,['plant']))
 return out,pd.DataFrame({'sample_id':df.index,'plant':groups,'distance_to_centroid':z})
dispersion_result,dispersion_values=dispersion(main)
print('dispersion',dispersion_result)
# Invariance to orthogonal rotation and restricted-permutation determinism.
ty2=rng.normal(size=(60,3)); rot=la.qr(rng.normal(size=(3,3)))[0]
a=fl_test(ty2,tx0,tx1,np.repeat(['a','b'],30),nperm=99)
b=fl_test(ty2@rot,tx0,tx1,np.repeat(['a','b'],30),nperm=99)
assert np.isclose(a['pseudo_F'],b['pseudo_F']) and a['p']==b['p']
validation={'univariate_F_identity':True,'fixed_seed_reproducibility':True,'orthogonal_invariance':True,'restricted_permutation_reproducibility':True,'distance_reconstruction_passed':True}


# %% Cell 7
cc=main.dropna(subset=['age','sex','bmi','antibiotic']).copy()
us=main.loc[main.country.eq('United States of America')].copy()
states=us.state.value_counts(); supported_states=states[states>=20].index.tolist()
regional=us.loc[us.state.isin(supported_states)].copy()
for frame in [main,cc,us,regional]:
 assert frame.plant.value_counts().min()>=20 and frame.plant.nunique()>=2
model('Complete cases unadjusted',cc,[],family='sensitivity')
model('Complete cases country-adjusted',cc,['country'],stratum='country',family='sensitivity')
model('Complete cases fully adjusted',cc,['country','age','sex','bmi','antibiotic'],stratum='country',family='sensitivity')
model('US only',us,[],family='sensitivity')
model('Eight states unadjusted',regional,[],family='sensitivity')
model('Eight states state-adjusted',regional,['state'],stratum='state',family='sensitivity')
response=cohort.loc[cohort.country.ne('Missing')].copy()
response['responded']=response.plant.notna().astype(int).astype(str)
assert response.responded.value_counts().min()>=20
model('Plant response status unadjusted',response,[],term='responded',family='missingness')
model('Plant response status country-adjusted',response,['country'],term='responded',stratum='country',family='missingness')
results_df=pd.DataFrame(results)
results_df['p_holm']=np.nan
for family in ['sensitivity','missingness']:
 idx=results_df.index[results_df.family.eq(family)]; vals=results_df.loc[idx,'p'].to_numpy(); order=np.argsort(vals)
 adj=np.minimum(1,np.maximum.accumulate(vals[order]*(len(vals)-np.arange(len(vals)))))
 results_df.loc[idx[order],'p_holm']=adj
response_table=cohort.assign(responded=cohort.plant.notna()).groupby('country').responded.agg(['size','sum','mean']).rename(columns={'size':'n','sum':'respondents','mean':'response_fraction'})
missing_table=pd.DataFrame({col:cohort[col].isna() for col in ['plant','state','sex','bmi','antibiotic']})
missing_table['country']=cohort.country.eq('Missing')
missing_summary=missing_table.agg(['sum','mean']).T.rename(columns={'sum':'missing_n','mean':'missing_fraction'})
complete_descriptive=main.assign(complete=main.index.isin(cc.index)).groupby('complete').agg(n=('plant','size'),mean_age=('age','mean'),median_depth=('depth','median'))
print(results_df.to_string(index=False))
print('response',response_table.to_dict('index'))
print('complete descriptive',complete_descriptive.to_dict('index'))
print('dispersion means',dispersion_values.groupby('plant').distance_to_centroid.agg(['size','mean']).to_dict('index'))


# %% Cell 8
import matplotlib.pyplot as plt
import os
out=pathlib.Path(os.environ.get('AG_OUTPUT_DIR','ag_results'))
out.mkdir(exist_ok=False)
results_df.to_csv(out/'model_results.csv',index=False)
pd.DataFrame(flow,columns=['stage','n']).to_csv(out/'cohort_flow.csv',index=False)
cohort[['sample_id','participant','country','state','plant','age','sex','bmi','antibiotic','depth','row']].to_csv(out/'cohort.csv',index=False)
pd.DataFrame({'feature_id':features[feat_mask]}).to_csv(out/'retained_features.csv',index=False)
pd.crosstab(cohort.country,cohort.plant).reindex(columns=LEVELS,fill_value=0).to_csv(out/'country_plant_support.csv')
response_table.to_csv(out/'response_by_country.csv')
missing_summary.to_csv(out/'questionnaire_missingness.csv')
complete_descriptive.to_csv(out/'complete_case_descriptives.csv')
dispersion_values.to_csv(out/'dispersion_values.csv',index=False)
pd.DataFrame([dispersion_result]).to_csv(out/'dispersion_test.csv',index=False)
pd.DataFrame(Y[:,:2],index=cohort.index,columns=['PCoA1','PCoA2']).to_csv(out/'ordination_display.csv',index_label='sample_id')
pd.DataFrame({'original':original_eigen,'corrected':ev}).to_csv(out/'geometry_eigenvalues.csv',index=False)
pd.DataFrame(manifest).to_csv(out/'download_manifest.csv',index=False)
versions={p:importlib.metadata.version(p) for p in ['numpy','pandas','scipy','matplotlib','h5py','threadpoolctl']}
versions['python']=sys.version
with (out/'validation.json').open('w') as f: json.dump({'checks':validation,'geometry':geometry,'versions':versions},f,indent=2)
field_dictionary={'TYPES_OF_PLANTS':'Reported plant-type category; exact historical recall interval not independently verified. Five labels preserved, unvalidated values missing.','HOST_SUBJECT_ID':'Numeric subject suffix; optional numeric study prefix removed for conservative deduplication.','BODY_SITE':'Select UBERON:feces.','AGE':'Numeric >=18 with AGE_UNIT years.','COUNTRY':'Recorded GAZ country stripped of prefix; unknown missing.','STATE':'Recorded region/state code; missing tokens excluded.','SEX':'female/male/other; otherwise missing.','BMI':'Positive finite numeric; no further plausibility exclusions.','ANTIBIOTIC_SELECT':'Five named time-since-antibiotic-use categories; otherwise missing.'}
pd.Series(field_dictionary,name='interpretation').to_csv(out/'data_dictionary.csv',index_label='source_field')
colors=['#332288','#117733','#0072B2','#CC6677','#882255']; markers=['o','s','^','D','v']
with plt.rc_context({'font.size':10,'axes.spines.top':False,'axes.spines.right':False,'pdf.fonttype':42,'savefig.facecolor':'white'}):
 fig,ax=plt.subplots(figsize=(8,6),layout='constrained')
 for g,color,marker in zip(LEVELS,colors,markers):
  ids=main.loc[main.plant.eq(g),'row'].to_numpy()
  ax.scatter(Y[ids,0],Y[ids,1],c=color,marker=marker,s=20,alpha=.55,label=f'{g} (n={len(ids)})')
 ax.set(xlabel=f'PCoA 1 ({100*geometry["axis1_fraction"]:.2f}% corrected inertia)',ylabel=f'PCoA 2 ({100*geometry["axis2_fraction"]:.2f}% corrected inertia)',title='Reported plant types and microbial composition')
 ax.legend(title='Reported plant-type category',fontsize=9,frameon=False)
 fig.savefig(out/'ordination.png',dpi=200); fig.savefig(out/'ordination.pdf'); plt.close(fig)
 fig,axs=plt.subplots(1,2,figsize=(13,5.5),layout='constrained',gridspec_kw={'width_ratios':[1.6,1]})
 shown=results_df.iloc[:8].copy(); yy=np.arange(len(shown))
 axs[0].scatter(100*shown.incremental_R2,yy,c='#0072B2',s=45)
 axs[0].set_yticks(yy,[f'{r.model} (n={r.n})' for r in shown.itertuples()]); axs[0].invert_yaxis()
 axs[0].set(xlim=(0,1.7),xlabel='Incremental explained variation (%)',title='A. Plant-term effect estimates')
 axs[0].grid(axis='x',alpha=.2)
 ms=missing_summary.loc[['plant','country','state','sex','bmi','antibiotic']]
 axs[1].barh(np.arange(len(ms)),ms.missing_n,color='#666666')
 axs[1].set_yticks(np.arange(len(ms)),['Plant types','Country','State','Sex','BMI','Antibiotics']); axs[1].invert_yaxis()
 axs[1].set(xlim=(0,65),xlabel='Missing responses (of 600)',title='B. Questionnaire completeness')
 for j,v in enumerate(ms.missing_n): axs[1].text(v+1,j,str(int(v)),va='center')
 fig.savefig(out/'sensitivity_missingness.png',dpi=200); fig.savefig(out/'sensitivity_missingness.pdf'); plt.close(fig)
from PIL import Image
for file in ['ordination.png','sensitivity_missingness.png']:
 with Image.open(out/file) as im: print('figure',file,im.size,im.info.get('dpi'))
print('OUTPUTS',str(out),len(list(out.iterdir())),'versions',versions)
