"""MAOMAO: raw clinical-event JSON -> checkpoint-compatible tensors -> predictions."""
import json,math
from pathlib import Path
import numpy as np
import torch
from safetensors.torch import load_file
from model import EventMAOMAO, dual_timescale_expected_wait

MODEL_ID='luan0519/MAOMAO'
FILES=['config.json','vocabulary.json','calibrations.json','model.safetensors']

def assets(directory=None):
    if directory:return Path(directory)
    from huggingface_hub import snapshot_download
    return Path(snapshot_download(MODEL_ID,allow_patterns=FILES))

class Predictor:
    def __init__(self,directory=None,device='cpu'):
        self.path=assets(directory);self.device=device
        self.vocab=json.loads((self.path/'vocabulary.json').read_text())
        cfg=json.loads((self.path/'config.json').read_text())
        self.model=EventMAOMAO(**cfg,outcome_family_ids=torch.tensor(self.vocab['outcome_to_family'])).eval()
        self.model.load_state_dict(load_file(str(self.path/'model.safetensors')))
        self.model.to(device)
        self.calibrations=json.loads((self.path/'calibrations.json').read_text())

    def encode(self,payload,device=None):
        if isinstance(payload,str):payload=json.loads(payload)
        if not isinstance(payload,dict):raise ValueError('Input must be a JSON object / 输入必须是 JSON 对象')
        rows=payload.get('events',[])
        if not isinstance(rows,list) or not 1<=len(rows)<=10000:raise ValueError('events: 1–10000 rows / 事件行数须为 1–10000')
        vocab=self.vocab['token_vocabulary'];outcome={f'event:{n}':i for i,n in enumerate(self.vocab['outcome_vocabulary'])}
        records=[dict(token='<BOS>',time_min=0,value=None)]
        for r in rows:
            if r.get('token') not in vocab or r['token'] in ['<PAD>','<MASK>','<EPISODE_END>']:raise ValueError(f"Unsupported input token / 不支持的事件: {r.get('token')}")
            t=float(r['time_min']);v=r.get('value')
            if not math.isfinite(t) or not 0<=t<=43200:raise ValueError('time_min must be finite and within 0–43200 / 时间须为 0–43200 分钟')
            if v is not None and not math.isfinite(float(v)):raise ValueError('value must be finite / 数值须有限')
            records.append(dict(token=r['token'],time_min=t,value=None if v is None else float(v)))
        records=sorted(records,key=lambda r:r['time_min'])
        s=payload.get('static',{});ranges={'age_at_operation':(0,120),'male':(0,1),'asa':(0,6),'emergency':(0,1),'weight_kg':(0,300),'height_cm':(0,250),'bmi':(0,80)}
        for k,(lo,hi) in ranges.items():
            v=float(s.get(k,0) or 0)
            if not math.isfinite(v) or not lo<=v<=hi:raise ValueError(f'{k}: range {lo}–{hi}')
        weight=float(s.get('weight_kg',0) or 0);height=float(s.get('height_cm',0) or 0)
        bmi=float(s.get('bmi',0) or (weight/(height/100)**2 if height>0 else 0));bmi=min(80,bmi)
        static=[float(s.get('age_at_operation',0) or 0)/100,float(s.get('male',0) or 0),bmi/40,float(s.get('asa',0) or 0)/6,float(s.get('emergency',0) or 0),weight/150,height/200]
        phase_map={'event:or_entry':1,'event:anesthesia_start':2,'event:surgery_start':3,'event:surgery_end':4,'event:or_exit':5,'event:icu_transfer':6,'event:icu_discharge':7}
        phase=0;last={};left=0;phases=[];observations=[];fams=self.vocab['outcome_to_family'];count=np.zeros(63,dtype=np.float32)
        start=max(0,len(records)-256)
        for i,r in enumerate(records):
            phase=phase_map.get(r['token'],phase);phases.append(phase)
            while r['time_min']-records[left]['time_min']>360:left+=1
            obs=[math.log1p(r['time_min']-records[left]['time_min'])/math.log1p(360),math.log1p(i-left)/math.log1p(128),0.]
            oi=outcome.get(r['token'])
            if oi is not None:
                fam=fams[oi]
                if fam in last:obs[2]=math.log1p(r['time_min']-last[fam])/math.log1p(43200)
                last[fam]=r['time_min']
                if i<start:count[fam]+=1
            observations.append(obs)
        selected=records[start:];times=[r['time_min'] for r in selected]
        def kind(token):
            if token=='<CLOCK>':return 7
            if token.startswith('event:'):return 6
            if token.startswith(('med:','med_signal:')):return 4
            if token.startswith('diagnosis:'):return 5
            if token.startswith(('procedure:','phase_summary:','context:department:','context:asa:','context:antype:')):return 3
            if token.startswith(('context:','summary:')):return 2
            return 1
        batch={'token_id':[vocab[r['token']] for r in selected],'token_kind':[kind(r['token']) for r in selected],
            'time_min':times,'gap_min':[0]+[times[i]-times[i-1] for i in range(1,len(times))],
            'value':[r['value'] or 0 for r in selected],'has_value':[float(r['value'] is not None) for r in selected],
            'phase_id':phases[start:],'observation_features':observations[start:],'attention_mask':[True]*len(selected),
            'history_family_counts':np.log1p(count).tolist(),'static':static}
        integer={'token_id','token_kind','phase_id'}
        return {k:torch.tensor(v,dtype=torch.long if k in integer else torch.bool if k=='attention_mask' else torch.float32,device=device or self.device).unsqueeze(0) for k,v in batch.items()}

    @torch.inference_mode()
    def predict(self,payload,calibration='None',custom=None,top_k=10):
        batch=self.encode(payload);output=self.model(batch)
        logits=output.logits[0,-1].float();raw=torch.softmax(logits,-1)
        fit={'temperature':1.,'bias':[0.]*len(logits)}
        if calibration=='Custom':fit=json.loads(custom) if isinstance(custom,str) else custom
        elif calibration!='None':fit=self.calibrations[calibration]
        t=float(fit['temperature']);bias=torch.as_tensor(fit.get('bias',[0.]*len(logits)),device=self.device)
        if not math.isfinite(t) or t<=0 or bias.shape!=logits.shape or not torch.isfinite(bias).all():raise ValueError('Invalid temperature or 210-element bias / 温度或 210 项偏置无效')
        corrected=torch.softmax((logits+bias)/t,-1)
        wait=dual_timescale_expected_wait(output.fine_hazard_logits[0,-1],output.long_hazard_logits[0,-1],output.tail_mu[0,-1],output.tail_log_sigma[0,-1])
        horizon=torch.sigmoid(output.trajectory_logits[0,-1]);names=self.vocab['outcome_vocabulary'];families=self.vocab['outcome_family_vocabulary'];fids=self.vocab['outcome_to_family']
        ids=torch.argsort(corrected,descending=True)[:int(top_k)].tolist()
        result=[]
        for rank,i in enumerate(ids,1):result.append({'rank':rank,'event':names[i],'family':families[fids[i]],'raw':float(raw[i]),'calibrated':float(corrected[i]),'wait_hours':float(wait[i]),'risk_1h_raw':float(horizon[0,i]),'risk_6h_raw':float(horizon[1,i]),'risk_24h_raw':float(horizon[2,i])})
        return {'calibration':calibration,'temperature':t,'input_events':len(batch['token_id'][0]),'next_event':result,'event_logits':logits.cpu().tolist(),'query_time_min':float(batch['time_min'][0,-1])}

if __name__=='__main__':
    import argparse
    p=argparse.ArgumentParser();p.add_argument('input');p.add_argument('--calibration',default='None');p.add_argument('--model-dir');a=p.parse_args()
    print(json.dumps(Predictor(a.model_dir).predict(Path(a.input).read_text(),a.calibration),indent=2))