#!/usr/bin/env python3
"""Position-specific barcode effects after removing each promoter's mean.

Gene holdouts evaluate within-gene residual prediction only. Centering test
responses is appropriate for that diagnostic, not prediction of new genes.
"""
import argparse
import hashlib
import json
import os
from pathlib import Path
import numpy as np
import pandas as pd

ROOT=Path(__file__).resolve().parents[1]
CACHE=Path(os.environ.get('BIO_DISCOVERY_CACHE','/mnt/data/research/bioinformatics-discovery/cache'))/'jores2021'

def read_counts(path,column):
    return pd.read_csv(path,sep=r'\s+',header=None,names=[column,'barcode'],compression='gzip')

def load(sp='At',system='leaf',rep=1,cutoff=5):
    roman={1:'I',2:'II'}
    sub=pd.read_csv(CACHE/f'analysis__subassembly__subassembly_{sp}.tsv',sep='\t')
    assert not sub.barcode.duplicated().any()
    prefix=f'data__barcode_counts__{system}__{sp}_Rep{rep}__barcodes_pPSup_{sp}PRO_35SEnh_{roman[rep]}'
    input_prefix=prefix
    if system=='proto' and sp=='At' and rep==2:
        input_prefix=f'data__barcode_counts__proto__At_Rep1__barcodes_pPSup_AtPRO_35SEnh_I'
    data=read_counts(CACHE/f'{prefix}_dark.count.gz','rna').merge(
        read_counts(CACHE/f'{input_prefix}_inp.count.gz','dna'),on='barcode',validate='one_to_one')
    data=data.merge(sub,on='barcode',validate='one_to_one')
    data=data[(data.variant=='WT')&data.FL&(data.gene!='35Spr')&(data.rna>=cutoff)&(data.dna>=cutoff)].copy()
    data=data[data.barcode.str.fullmatch('[ACGT]{12}')]
    data['y']=np.log2(data.rna/data.dna)
    data['group_size']=data.groupby('gene').barcode.transform('size')
    return data[data.group_size>=2].reset_index(drop=True)

def design(data):
    columns={}
    for pos in range(12):
        for base in 'CGT':
            columns[f'pos{pos+1}_{base}']=(data.barcode.str[pos]==base).astype(float)
    return pd.DataFrame(columns)

def main():
    parser=argparse.ArgumentParser();parser.add_argument('--sp',default='At');parser.add_argument('--system',default='leaf');parser.add_argument('--rep',type=int,default=1)
    args=parser.parse_args();data=load(args.sp,args.system,args.rep)
    raw=design(data)
    x=raw-raw.groupby(data.gene).transform('mean')
    y=data.y-data.groupby('gene').y.transform('mean')
    valid=data.gene.map(lambda g:int(hashlib.sha256(f'20260919:{g}'.encode()).hexdigest()[:8],16)%5==0).to_numpy()
    beta=np.linalg.lstsq(x.to_numpy()[~valid],y.to_numpy()[~valid],rcond=None)[0]
    pred=x.to_numpy()@beta
    metrics=dict(species=args.sp,system=args.system,replicate=args.rep,barcodes=len(data),promoters=data.gene.nunique(),
        train_promoters=data.loc[~valid,'gene'].nunique(),test_promoters=data.loc[valid,'gene'].nunique(),
        heldout_within_promoter_r2=float(1-np.sum((y.to_numpy()[valid]-pred[valid])**2)/np.sum(y.to_numpy()[valid]**2)),
        first_base_ACG_span=float(np.ptp([0,beta[0],beta[1]])),
        first_base_counts=data.barcode.str[0].value_counts().to_dict(),
        note='Observational OLS after within-promoter centering; >=5 counts in both DNA and RNA. No mechanistic attribution. ACG first-base span excludes the design-disallowed T.')
    name=f'010_{args.sp}_{args.system}_rep{args.rep}'
    pd.DataFrame(dict(feature=raw.columns,coefficient=beta)).to_csv(ROOT/f'results/{name}_coefficients.csv',index=False)
    (ROOT/f'results/{name}_metrics.json').write_text(json.dumps(metrics,indent=2)+'\n')
    # Cache inputs for validation, without committing raw barcode sequences.
    data.to_pickle(CACHE/f'{name}_data.pkl')
    print(json.dumps(metrics,indent=2))
    print(pd.DataFrame(dict(feature=raw.columns,coefficient=beta)).head(12).to_string(index=False))

if __name__=='__main__':main()