#!/usr/bin/env python3
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor
import hashlib,time,json,requests
import pandas as pd
import numpy as np
ROOT=Path(__file__).resolve().parents[1]; DAT=ROOT/'data'; OUT=ROOT/'output'
DAT.mkdir(exist_ok=True);OUT.mkdir(exist_ok=True)
manifest=[]
def download(job):
 year,name,ext,url=job;p=DAT/(name+'.'+ext)
 if not p.exists():
  for k in range(8):
   try:
    r=requests.get(url,timeout=90)
    good=r.status_code==200 and (r.content.startswith(b'HEADER') if ext=='XPT' else b'<html' not in r.content[:200].lower())
    if good:p.write_bytes(r.content);break
   except requests.RequestException:pass
  else:return dict(file=p.name,url=url,status='FAILED',bytes=0,sha256='',download_utc='')
 return dict(file=p.name,url=url,status='ok',bytes=p.stat().st_size,sha256=hashlib.sha256(p.read_bytes()).hexdigest(),download_utc=time.strftime('%Y-%m-%dT%H:%M:%SZ',time.gmtime(p.stat().st_mtime)))
def n(x):return pd.to_numeric(x,errors='coerce')
def main():
 jobs=[];spec=[]
 for i,y in enumerate(range(1999,2019,2)):
  suf='' if i==0 else '_'+chr(65+i)
  glu=['LAB10AM','L10AM_B','L10AM_C'][i] if i<3 else 'GLU'+suf
  tri=['LAB13AM','L13AM_B','L13AM_C'][i] if i<3 else 'TRIGLY'+suf
  files=['DEMO'+suf,glu,tri,'BMX'+suf,'SMQ'+suf,'DIQ'+suf,'BPQ'+suf,'MCQ'+suf,'ALQ'+suf]
  spec.append((y,files))
  for f in files:jobs.append((y,f,'XPT',f'https://wwwn.cdc.gov/Nchs/Data/Nhanes/Public/{y}/DataFiles/{f}.xpt'))
  f=f'NHANES_{y}_{y+1}_MORT_2019_PUBLIC';jobs.append((y,f,'dat','https://ftp.cdc.gov/pub/Health_Statistics/NCHS/datalinkage/linked_mortality/'+f+'.dat'))
 with ThreadPoolExecutor(8) as ex:manifest=list(ex.map(download,jobs))
 pd.DataFrame(manifest).to_csv(OUT/'data_manifest.csv',index=False)
 allx=[];mapping=[]
 for y,files in spec:
  frames=[]
  for f in files:
   p=DAT/(f+'.XPT')
   if not p.exists():raise RuntimeError('Missing '+str(p))
   d=pd.read_sas(p,format='xport'); keep=['SEQN','RIDAGEYR','RIAGENDR','RIDRETH1','DMDEDUC2','INDFMPIR','RIDEXPRG','SDMVPSU','SDMVSTRA','WTSAF4YR','WTSAF2YR','LBXGLU','LBXTR','BMXBMI','SMQ020','SMQ040','DIQ010','BPQ020','MCQ160B','MCQ160C','MCQ160D','MCQ160E','MCQ160F','MCQ220','ALQ101']; frames.append(d[[c for c in keep if c in d]]);mapping.extend([dict(cycle=f'{y}-{y+1}',file=f,variable=c) for c in d.columns])
  x=frames[0]
  for d in frames[1:]:x=x.merge(d.drop(columns=[c for c in d if c!='SEQN' and c in x]),on='SEQN',how='left',validate='1:1')
  m=pd.read_fwf(DAT/f'NHANES_{y}_{y+1}_MORT_2019_PUBLIC.dat',colspecs=[(0,6),(14,15),(15,16),(16,19),(19,20),(20,21),(42,45),(45,48)],names=['SEQN','ELIGSTAT','MORTSTAT','UCOD_LEADING','DIABETES','HYPERTEN','PERMTH_INT','PERMTH_EXM'],na_values=['.'])
  x=x.merge(m,on='SEQN',how='left',validate='1:1');x['cycle']=f'{y}-{y+1}'
  x['weight']=x['WTSAF4YR']*.2 if y<2003 else x['WTSAF2YR']*.1
  allx.append(x)
 pd.DataFrame(mapping).to_csv(OUT/'variable_mapping.csv',index=False)
 x=pd.concat(allx,ignore_index=True);x.to_csv(OUT/'merged_all.csv',index=False)
 def col(c):return n(x[c]) if c in x else pd.Series(np.nan,index=x.index)
 x['age']=col('RIDAGEYR');x['sex']=col('RIAGENDR');x['race']=col('RIDRETH1');x['education']=col('DMDEDUC2').where(col('DMDEDUC2').between(1,5));x['pir']=col('INDFMPIR');x['bmi']=col('BMXBMI')
 x['glucose']=col('LBXGLU');x['triglycerides']=col('LBXTR');x['tyg']=np.log(x.glucose*x.triglycerides/2)
 smoke=col('SMQ020');now=col('SMQ040');x['smoking']=np.select([smoke.eq(2),smoke.eq(1)&now.eq(3),smoke.eq(1)&now.isin([1,2])],[0,1,2],default=np.nan)
 # Comorbidities combine self report with fasting glucose for diabetes.
 x['diabetes']=np.where(col('DIQ010').eq(1)|x.glucose.ge(126),1,np.where(col('DIQ010').isin([2,3])&x.glucose.notna(),0,np.nan))
 x['hypertension']=np.where(col('BPQ020').eq(1),1,np.where(col('BPQ020').eq(2),0,np.nan))
 cv=x[['MCQ160B','MCQ160C','MCQ160D','MCQ160E','MCQ160F']].apply(n);x['cvd']=np.where(cv.eq(1).any(axis=1),1,np.where(cv.eq(2).all(axis=1),0,np.nan))
 x['cancer']=np.where(col('MCQ220').eq(1),1,np.where(col('MCQ220').eq(2),0,np.nan))
 x['alcohol']=np.where(col('ALQ101').eq(1),1,np.where(col('ALQ101').eq(2),0,np.nan))
 x['time']=col('PERMTH_EXM')/12;x['death']=col('MORTSTAT');x['cvdeath']=((x.death==1)&col('UCOD_LEADING').isin([1,5])).astype(int)
 flow=[]
 def step(label,mask):
  nonlocal x
  before=len(x);x=x[mask].copy();flow.append(dict(step=label,before_n=before,excluded_n=before-len(x),remaining_n=len(x)))
 flow.append(dict(step='All NHANES participants',before_n=len(x),excluded_n=0,remaining_n=len(x)))
 step('Age >=20',x.age.ge(20));step('Not pregnant (RIDEXPRG !=1)',~col('RIDEXPRG').eq(1))
 step('Positive glucose and triglycerides',x.glucose.gt(0)&x.triglycerides.gt(0));step('Positive fasting weight',x.weight.gt(0));step('Mortality eligible',x.ELIGSTAT.eq(1));step('Valid mortality and positive exam follow-up',x.death.isin([0,1])&x.time.gt(0))
 x['weight_20yr']=x['weight'];x['cvd_death']=x['cvdeath'];x['cvd_history']=x['cvd'];x.to_csv(OUT/'analysis.csv',index=False);pd.DataFrame(flow).to_csv(OUT/'flow_counts.csv',index=False)
 pd.DataFrame([dict(n=len(x),deaths=int(x.death.sum()),cv_deaths=int(x.cvdeath.sum()),tyg_min=x.tyg.min(),tyg_max=x.tyg.max())]).to_csv(OUT/'data_summary.csv',index=False)
 print(pd.DataFrame(flow).to_string(index=False));print(x.groupby('cycle').agg(n=('SEQN','size'),deaths=('death','sum'),cv_deaths=('cvdeath','sum')).to_string())
if __name__=='__main__':main()