"""Aggregate-only life table calculator; Python standard library.
python3 scripts/calculate.py data/demo_bundle.json --end-age 65 --reference C
"""
import argparse,json,math
from datetime import date
from pathlib import Path

class InputError(ValueError): pass

def number(value,label):
 try: n=float(value)
 except (TypeError,ValueError): raise InputError(f'{label}: 숫자가 필요합니다.')
 if not math.isfinite(n):raise InputError(f'{label}: 유한한 숫자가 필요합니다.')
 return n

def indexed(rows,key,label):
 out={}
 for r in rows:
  k=number(r.get(key),label+' '+key)
  if k!=int(k):raise InputError(label+': 정수 연령이 필요합니다.')
  k=int(k)
  if k in out:raise InputError(f'{label}: 연령 {k} 중복입니다.')
  out[k]=r
 return out

def quant_ci(value,se):return [value-1.96*se,value+1.96*se]

def calculate(bundle,end_age=60,reference='C'):
 if end_age not in (60,64,65):raise InputError('지원 기준연령은 60, 64, 65세입니다.')
 meta=bundle.get('metadata',{})
 for field in ['analysis_id','design','classification_rule','observation_start','observation_end','age_method','sex','classifier','data_kind','standard_population','reference_source']:
  if not meta.get(field) or meta.get(field)=='미입력':raise InputError(f'메타정보 {field}가 필요합니다.')
 try:
  if date.fromisoformat(meta['observation_start'])>date.fromisoformat(meta['observation_end']):raise ValueError()
 except (ValueError,TypeError):raise InputError('관찰 시작·종료일은 순서가 맞는 YYYY-MM-DD 형식이어야 합니다.')
 if meta['data_kind'] not in ['synthetic','real_aggregate']:raise InputError('data_kind는 synthetic 또는 real_aggregate입니다.')
 if meta['design'] not in ['cohort_10y','age40','cross_section_1y','stayer_1y']:raise InputError('분석설계 코드가 올바르지 않습니다.')
 last=number(meta.get('observed_last_age'),'observed_last_age')
 if last not in (59,63):raise InputError('현재 버전의 관측 마지막 연령은 59 또는 63세입니다.')
 last=int(last)
 if meta['design']=='cross_section_1y' and last!=59:raise InputError('현재 1년 단면 규약은 40–59세 관측 자료입니다.')
 if last==59 and end_age!=60:raise InputError('1년 단면 예제는 관측된 40–59세만 계산합니다. 5개 연령 외삽은 이 데모에서 제공하지 않습니다.')
 rows=bundle.get('aggregate',[])
 if not rows or len(rows)>100000:raise InputError('집계표는 1–100,000행이어야 합니다.')
 groups={};labels={}
 for r in rows:
  for k in ['analysis_id','sex','classifier']:
   if str(r.get(k))!=str(meta[k]):raise InputError(f'{k}가 메타정보와 다릅니다. 서로 다른 설계를 섞을 수 없습니다.')
  code=str(r.get('group_code',''))
  if not code:raise InputError('group_code가 필요합니다.')
  label=str(r.get('group_label',''))
  if not label:raise InputError('group_label이 필요합니다.')
  if code in labels and label!=labels[code]:raise InputError('동일 집단의 이름이 일치하지 않습니다.')
  labels[code]=label;groups.setdefault(code,[]).append(r)
 if reference not in groups:raise InputError('기준집단이 집계표에 없습니다.')
 weights=indexed(bundle.get('standard_population',[]),'age','표준인구')
 refs=indexed(bundle.get('reference_mortality',[]),'age','외삽 기준')
 obs_end=min(end_age,last+1)
 ws={}
 for age in range(40,obs_end):
  if age not in weights:raise InputError(f'표준인구의 {age}세 값이 없습니다.')
  w=number(weights[age].get('weight'),'표준인구 weight')
  if w<0:raise InputError('표준인구 가중치는 음수일 수 없습니다.')
  ws[age]=w
 if sum(ws.values())<=0:raise InputError('표준인구 가중치 합이 0입니다.')
 norm=sum(ws.values());ws={a:v/norm for a,v in ws.items()}
 output=[]
 for code,raw in groups.items():
  cells=indexed(raw,'age',code)
  if any(a<40 or a>last for a in cells):raise InputError(f'{code}: 관측 연령 범위 밖의 행이 있습니다.')
  parsed={}
  for age in range(40,obs_end):
   if age not in cells:raise InputError(f'{code}: {age}세 행이 없습니다. 누락값을 사망 0건으로 대체할 수 없습니다.')
   r=cells[age]
   if r.get('cell_status')!='observed':raise InputError(f'{code} {age}세: 비공개/결측 셀 때문에 계산할 수 없습니다.')
   py=number(r.get('person_years'),'person_years');d=number(r.get('deaths'),'deaths')
   if py<=0:raise InputError(f'{code} {age}세: 인년은 양수여야 합니다.')
   if d<0 or d!=int(d):raise InputError('사망 수는 0 이상의 정수여야 합니다.')
   if d/py>2:raise InputError('사망확률이 1을 넘는 입력입니다. 사망 건수 및 인년 단위 확인이 필요합니다.')
   parsed[age]=(py,d,r)
  mx={a:d/py for a,(py,d,r) in parsed.items()};extra=0;k=None;dk=0
  if end_age==65:
   if last!=63:raise InputError('64세 1년 외삽은 63세까지 관측된 집계표에서만 지원합니다.')
   for age in [61,62,63,64]:
    if age not in refs:raise InputError(f'외삽 기준 사망률 {age}세가 없습니다.')
    if number(refs[age].get('mx'),'기준 mx')<=0:raise InputError('외삽 기준 사망률은 양수여야 합니다.')
   dk=sum(parsed[a][1] for a in (61,62,63))
   expected=sum(parsed[a][0]*number(refs[a]['mx'],'기준 mx') for a in (61,62,63))
   k=dk/expected;mx[64]=k*number(refs[64]['mx'],'기준 mx');extra=1
   if mx[64]>2:raise InputError('외삽 사망확률이 1을 넘습니다.')
  life=[];survival=1.0;partial=0.0
  for age in range(40,end_age):
   m=mx[age];q=m/(1+.5*m);L=survival*(1-.5*q)
   life.append({'age':age,'mx':m,'qx':q,'lx':survival,'Lx':L,'source':'extrapolated' if age==64 else 'observed'})
   partial+=L;survival*=1-q
  remaining=0.0;variance=0.0
  for cell in reversed(life):
   age=cell['age'];q=cell['qx'];l=cell['lx'];next_l=l*(1-q)
   next_e=remaining/next_l if next_l>0 else 0
   if age in parsed:
    d=parsed[age][1];vq=q*q*(1-q)/d if d>0 else 0
    variance+=l*l*(.5+next_e)**2*vq
   remaining+=cell['Lx']
  if extra:
   c=life[-1];m=mx[64]
   # Delta variance from k alone (same approximation as thesis main calculation).
   # Covariance with ages 61–63 intentionally not added; reported as assumption.
   lost_extra=.5*c['lx']*c['qx']
   variance+=lost_extra**2/dk if dk>0 else 0
  dtotal=sum(v[1] for v in parsed.values());pytotal=sum(v[0] for v in parsed.values())
  loss=end_age-40-partial;se=math.sqrt(variance)
  ci=quant_ci(loss,se) if dtotal>=20 else None
  ytotal=0.0;yrate=0.0;yvar=0.0;has_ypll=True
  for age,(py,d,r) in parsed.items():
   a1=r.get('sum_age_at_death');a2=r.get('sum_age_at_death_sq')
   if (a1 in ['',None])!=(a2 in ['',None]):raise InputError('사망연령 합·제곱합은 함께 제공해야 합니다.')
   if a1 not in ['',None]:
    a1=number(a1,'사망연령 합');a2=number(a2,'사망연령 제곱합')
    if d==0 and (a1!=0 or a2!=0):raise InputError('사망 0건인 셀의 사망연령 합·제곱합은 0이어야 합니다.')
    if d>0:
     if not age*d-1e-7<=a1<=(age+1)*d+1e-7:raise InputError('사망연령 합이 해당 만 나이 구간과 맞지 않습니다.')
     if a2+1e-6<a1*a1/d or a2>(2*age+1)*a1-age*(age+1)*d+1e-6:raise InputError('사망연령 제곱합이 가능한 범위를 벗어났습니다.')
    y=end_age*d-a1;y2=end_age**2*d-2*end_age*a1+a2
   else:
    y=r.get(f'ypll_{end_age}_sum');y2=r.get(f'ypll_{end_age}_sq_sum')
    if y in ['',None] or y2 in ['',None]:has_ypll=False;continue
    y=number(y,'YPLL 합');y2=number(y2,'YPLL 제곱합')
    if d==0 and (y!=0 or y2!=0):raise InputError('사망 0건 셀의 YPLL은 0이어야 합니다.')
    if d and y2+1e-6<y*y/d:raise InputError('YPLL 제곱합이 합과 일치하지 않습니다.')
    if d and (y < d*(end_age-age-1-.01) or y > d*(end_age-age+.01) or y2 > d*(end_age-age+.01)**2):raise InputError('고정 기준연령 YPLL이 해당 만 나이 구간과 맞지 않습니다.')
   if y< -1e-6 or y2< -1e-6:raise InputError('YPLL 합·제곱합은 음수일 수 없습니다.')
   y=max(0,y);y2=max(0,y2);ytotal+=y;yrate+=ws[age]*y/py;yvar+=ws[age]**2*y2/py**2
  output.append({'group_code':code,'group_label':labels[code],'deaths':dtotal,'person_years':pytotal,'partial_life_expectancy':partial,'life_years_lost':loss,'se':se if dtotal>=20 else None,'ci_low':ci[0] if ci else None,'ci_high':ci[1] if ci else None,'inference_status':'available' if ci else 'fewer_than_20_deaths','death_probability':1-survival,'extrapolated_years':extra,'extrapolation_factor':k,'ypll_sum':ytotal if has_ypll else None,'ypll_rate':yrate*1e5 if has_ypll else None,'ypll_rate_se':math.sqrt(yvar)*1e5 if has_ypll and dtotal>=20 else None,'ypll_age_range':f'40–{obs_end-1}','life_table':life})
 ref=next(r for r in output if r['group_code']==reference)
 for r in output:
  r['difference']=r['life_years_lost']-ref['life_years_lost'];r['ratio']=r['life_years_lost']/ref['life_years_lost'] if ref['life_years_lost']>0 else None
  r['difference_ci_low']=r['difference_ci_high']=r['ratio_ci_low']=r['ratio_ci_high']=None
  r['p_raw']=r['p_holm']=None
  if r['group_code']==reference:continue
  if r['se'] is not None and ref['se'] is not None:
   dse=math.hypot(r['se'],ref['se']);r['difference_ci_low'],r['difference_ci_high']=quant_ci(r['difference'],dse)
   r['p_raw']=math.erfc(abs(r['difference'])/dse/math.sqrt(2)) if dse else (1.0 if r['difference']==0 else None)
   if r['ratio'] and r['life_years_lost']>0:
    logse=math.hypot(r['se']/r['life_years_lost'],ref['se']/ref['life_years_lost'])
    r['ratio_ci_low']=math.exp(math.log(r['ratio'])-1.96*logse);r['ratio_ci_high']=math.exp(math.log(r['ratio'])+1.96*logse)
 family=sorted([r for r in output if r['p_raw'] is not None],key=lambda r:r['p_raw']);previous=0
 for i,r in enumerate(family):
  previous=max(previous,min(1.0,r['p_raw']*(len(family)-i)));r['p_holm']=previous
 return {'metadata':meta,'settings':{'start_age':40,'end_age':end_age,'reference':reference},'method_version':'0.1.0','notes':['합성자료 기반 계산 결과. 논문의 실제 분석 결과와 무관.' if meta['data_kind']=='synthetic' else '사용자 입력 집계자료. 원자료 및 공개 권한의 별도 확인 필요.','차이·비율의 신뢰구간: 집단 간 독립 가정에 따른 근사값. 겸직 중복에 따른 공분산 미반영.','64세 신뢰구간: 기준 사망률을 고정하고 외삽계수의 변동만 반영한 근사값. 관측 연령과의 공분산 및 외삽 가정의 불확실성 미반영.','YPLL 산출 범위: 관측 사망만 포함. 외삽 사망 건수 미산출.','사망 20건 미만 집단: 점추정값만 제시. 공개용 최소 셀 기준과는 별개.'],'results':output}

if __name__=='__main__':
 parser=argparse.ArgumentParser();parser.add_argument('bundle',type=Path);parser.add_argument('--end-age',type=int,default=60);parser.add_argument('--reference',default='C');parser.add_argument('--output',type=Path);a=parser.parse_args()
 try: result=calculate(json.loads(a.bundle.read_text(encoding='utf-8-sig')),a.end_age,a.reference)
 except (InputError,KeyError,TypeError,json.JSONDecodeError) as e:parser.exit(2,f'입력 오류: {e}\n')
 text=json.dumps(result,ensure_ascii=False,indent=2,allow_nan=False)+'\n'
 if a.output:a.output.write_text(text,encoding='utf-8')
 else:print(text)