"""PanTHERIA retrospective benchmark. Run with Python, NumPy and SciPy.

Reads the unchanged publisher file, preserves row exclusions, and never searches
for a breakpoint. Family blocking is a sensitivity to dependence, not PGLS.
"""
from pathlib import Path
import argparse, csv, hashlib, json, math, platform
import numpy as np
from scipy.optimize import least_squares

HERE = Path(__file__).resolve().parent
MODELS = ('M0', 'M1_10kg', 'M2_quadratic', 'M1_20kg')
SEED = 20261020

def design(x, model):
    if model not in MODELS: raise ValueError('Unknown model: '+model)
    if model == 'M0': return np.column_stack([np.ones(len(x)), x])
    if model == 'M2_quadratic': return np.column_stack([np.ones(len(x)), x, x*x])
    hinge = 1.0 if model == 'M1_10kg' else math.log10(20)
    return np.column_stack([np.ones(len(x)), x, np.maximum(0, x-hinge)])

def load_data(path):
    with path.open(encoding='utf-8-sig', newline='') as f: raw = list(csv.DictReader(f, delimiter='\t'))
    selected, exclusions, seen = [], [], set()
    for index, row in enumerate(raw, 2):
        name = row['MSW05_Binomial'].strip()
        if name in seen: raise ValueError('Duplicate binomial requires explicit adjudication: '+name)
        seen.add(name)
        mass = float(row['5-2_BasalMetRateMass_g'])
        oxygen = float(row['18-1_BasalMetRate_mLO2hr'])
        family = row['MSW05_Family'].strip()
        reasons = []
        if mass == -999: reasons.append('missing BMR-associated mass')
        elif not math.isfinite(mass) or mass <= 0: reasons.append('nonpositive or invalid mass')
        if oxygen == -999: reasons.append('missing BMR')
        elif not math.isfinite(oxygen) or oxygen <= 0: reasons.append('nonpositive or invalid BMR')
        if not family or family == '-999': reasons.append('missing family')
        if reasons:
            exclusions.append(dict(source_line=index, binomial=name, reason='; '.join(reasons)))
        else:
            selected.append(dict(source_line=index, binomial=name, family=family,
                order=row['MSW05_Order'], mass_kg=mass/1000,
                bmr_mLO2_hr=oxygen, bmr_watt=oxygen*20.1/3600))
    return sorted(selected, key=lambda r:r['binomial']), exclusions, len(raw)

def split(rows, grouped=True):
    keys = [r['family'] if grouped else r['binomial'] for r in rows]
    groups = sorted(set(keys), key=lambda k:(np.median([math.log10(r['mass_kg']) for r,v in zip(rows, keys) if v==k]), k))
    test_keys = set()
    for block in np.array_split(np.array(groups, dtype=object), 10):
        block = sorted(block, key=lambda k:hashlib.sha256((str(SEED)+':'+str(k)).encode()).hexdigest())
        ntest = min(len(block)-1, max(1, int(math.floor(.3*len(block)+.5)))) if len(block)>=2 else 0
        test_keys.update(block[:ntest])
    test = np.array([i for i,k in enumerate(keys) if k in test_keys], dtype=int)
    train = np.array([i for i,k in enumerate(keys) if k not in test_keys], dtype=int)
    return train, test

def fit(x, y, model, idx):
    X=design(x[idx], model)
    if np.linalg.matrix_rank(X) != X.shape[1]: raise ValueError('Rank-deficient design for '+model)
    coef=np.linalg.lstsq(X, y[idx], rcond=None)[0]
    resid=y[idx]-X@coef; n=len(idx); p=X.shape[1]; k=p+1
    sse=float(resid@resid); var=max(sse/n,1e-30)
    ll=-.5*n*(math.log(2*math.pi*var)+1)
    aic=-2*ll+2*k
    return coef, dict(coefficients=coef.tolist(), n_train=n, parameters_including_variance=k,
        residual_sd_dex=math.sqrt(sse/(n-p)), training_log_likelihood=ll,
        training_aicc=aic+2*k*(k+1)/(n-k-1) if n>k+1 else None)

def evaluate(x,y,train,test):
    results={}; predictions={}
    for model in MODELS:
        coef, info=fit(x,y,model,train)
        pred=design(x[test],model)@coef; error=y[test]-pred
        info.update(test_rmse_dex=float(np.sqrt(np.mean(error**2))),test_mse_dex2=float(np.mean(error**2)),test_bias_dex=float(np.mean(error)))
        if model.startswith('M1'): info.update(lower_slope=float(coef[1]),upper_slope=float(coef[1]+coef[2]),slope_change=float(coef[2]))
        results[model]=info;predictions[model]=pred
    minimum=min(v['training_aicc'] for v in results.values()); weights=np.exp(-.5*np.array([v['training_aicc']-minimum for v in results.values()]));weights/=weights.sum()
    for (model, info),weight in zip(results.items(),weights): info.update(training_delta_aicc=info['training_aicc']-minimum,training_akaike_weight=float(weight))
    return results,predictions

def cluster_resample(rows,idx,rng):
    groups=sorted(set(rows[i]['family'] for i in idx))
    mapping={g:idx[[rows[i]['family']==g for i in idx]] for g in groups}
    return np.concatenate([mapping[g] for g in rng.choice(groups,len(groups),replace=True)])

def bootstraps(rows,x,y,train,test,reps):
    rng=np.random.default_rng(SEED); samples=[]; failed=0
    for _ in range(reps):
        tr=cluster_resample(rows,train,rng);te=cluster_resample(rows,test,rng)
        try:
            coef={m:fit(x,y,m,tr)[0] for m in MODELS}
        except ValueError:
            failed+=1;continue
        losses={m:float(np.mean((y[te]-design(x[te],m)@coef[m])**2)) for m in MODELS}
        samples.append([losses['M0']-losses['M1_10kg'],losses['M2_quadratic']-losses['M1_10kg'],float(coef['M1_10kg'][2]),1-math.sqrt(losses['M1_10kg']/losses['M0'])])
    if len(samples)<reps*.9: raise ValueError('Too many unidentified bootstrap samples')
    a=np.array(samples);keys=['M0_minus_M1_test_mse','M2_minus_M1_test_mse','M1_slope_change','M1_rmse_fraction_improvement']
    return dict(requested=reps,valid=len(samples),rank_deficient=failed,intervals={k:np.quantile(a[:,i],[.025,.975]).tolist() for i,k in enumerate(keys)})

def family_cv(rows,x,y):
    families=sorted(set(r['family'] for r in rows),key=lambda g:hashlib.sha256(('cv:'+str(SEED)+':'+g).encode()).hexdigest())
    allocation={g:i%5 for i,g in enumerate(families)};out=[]
    for fold in range(5):
        tr=np.array([i for i,r in enumerate(rows) if allocation[r['family']]!=fold]);te=np.array([i for i,r in enumerate(rows) if allocation[r['family']]==fold])
        res,_=evaluate(x,y,tr,te)
        out.append(dict(fold=fold,n_test=len(te),test_rmse_dex={m:res[m]['test_rmse_dex'] for m in MODELS}))
    pooled={m:math.sqrt(sum(f['n_test']*f['test_rmse_dex'][m]**2 for f in out)/len(rows)) for m in MODELS}
    return dict(folds=out,pooled_out_of_fold_rmse_dex=pooled,scope='Sensitivity across all families; does not create a new untouched test set.')

def additive_sensitivity(x,y,train,test):
    raw=10**y;scale=float(np.mean(raw[train]));out={}
    for m in MODELS:
        init,_=fit(x,y,m,train);X=design(x[train],m)
        f=least_squares(lambda b:(10**np.clip(X@b,-20,20)-raw[train])/scale,init,max_nfev=5000)
        pred=10**np.clip(design(x[test],m)@f.x,-20,20)
        out[m]=dict(converged=bool(f.success),coefficients=f.x.tolist(),test_rmse_watt=float(np.sqrt(np.mean((raw[test]-pred)**2))),test_rmse_dex=float(np.sqrt(np.mean((y[test]-np.log10(pred))**2))))
    return out

def write_csv(path,rows):
    if not rows:return
    with path.open('w',newline='') as f:
        w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)

def main():
    parser=argparse.ArgumentParser();parser.add_argument('--input',type=Path,default=HERE/'PanTHERIA_1-0_WR05_Aug2008.txt');parser.add_argument('--output',type=Path,default=HERE/'metabolic-output');parser.add_argument('--bootstrap',type=int,default=2000);args=parser.parse_args()
    if args.bootstrap<100:raise ValueError('At least 100 bootstrap draws required; release uses 2000')
    rows,excluded,nraw=load_data(args.input)
    if len(rows)<50: raise ValueError('Insufficient eligible species')
    x=np.array([math.log10(r['mass_kg']) for r in rows]);y=np.array([math.log10(r['bmr_watt']) for r in rows])
    train,test=split(rows,True);st,sq=split(rows,False)
    assert not ({rows[i]['family'] for i in train}&{rows[i]['family'] for i in test})
    results,predictions=evaluate(x,y,train,test);sensitivity,_=evaluate(x,y,st,sq)
    boot=bootstraps(rows,x,y,train,test,args.bootstrap)
    counts=dict(raw_rows=nraw,eligible_species=len(rows),excluded_rows=len(excluded),families=len(set(r['family'] for r in rows)),training_species=len(train),test_species=len(test),training_families=len(set(rows[i]['family'] for i in train)),test_families=len(set(rows[i]['family'] for i in test)),above_10kg_total=sum(r['mass_kg']>10 for r in rows),above_10kg_train=sum(rows[i]['mass_kg']>10 for i in train),above_10kg_test=sum(rows[i]['mass_kg']>10 for i in test),above_10kg_families=len(set(r['family'] for r in rows if r['mass_kg']>10)))
    plan=json.loads((HERE/'analysis-plan.json').read_text());coverage=all(counts[k]>=v for k,v in plan['coverage_gate'].items())
    improvement=1-results['M1_10kg']['test_rmse_dex']/results['M0']['test_rmse_dex'];ci=boot['intervals']
    screen=coverage and improvement>.05 and ci['M0_minus_M1_test_mse'][0]>0 and ci['M1_slope_change'][0]>0
    outputs=dict(status='Retrospective analysis of externally measured species summaries; not preregistered, independent replication or new NetPhysics measurements.',input_sha256=hashlib.sha256(args.input.read_bytes()).hexdigest(),analysis_plan_sha256=hashlib.sha256((HERE/'analysis-plan.json').read_bytes()).hexdigest(),counts=counts,primary_family_blocked=results,bootstrap=boot,coverage_gate_pass=coverage,hinge_vs_single_screen_pass=screen,hinge_rmse_fraction_improvement=improvement,species_split_sensitivity=sensitivity,five_fold_family_cv=family_cv(rows,x,y),additive_error_sensitivity=additive_sensitivity(x,y,train,test),limitations=plan['open_limits'],environment=dict(python=platform.python_version(),numpy=np.__version__))
    args.output.mkdir(exist_ok=True,parents=True)
    (args.output/'results.json').write_text(json.dumps(outputs,indent=2,allow_nan=False)+'\n')
    for i,r in enumerate(rows):r['primary_split']='train' if i in train else 'test';r['species_sensitivity_split']='train' if i in st else 'test'
    write_csv(args.output/'included-species.csv',rows);write_csv(args.output/'excluded-species.csv',excluded)
    preds=[]
    for j,i in enumerate(test):
        r=dict(binomial=rows[i]['binomial'],family=rows[i]['family'],mass_kg=rows[i]['mass_kg'],observed_log10_watt=float(y[i]))
        r.update({m+'_predicted_log10_watt':float(predictions[m][j]) for m in MODELS});preds.append(r)
    write_csv(args.output/'held-out-predictions.csv',preds)
    print(json.dumps(dict(counts=counts,coverage=coverage,screen=screen,test_rmse={m:r['test_rmse_dex'] for m,r in results.items()},bootstrap_intervals=ci),indent=2))

if __name__=='__main__':main()