#!/usr/bin/env python3
"""Entry 021: fixed chromosome-1 terminal-ATG test, no validation outcomes."""
import itertools,json,re
from pathlib import Path
import numpy as np
import pandas as pd
ROUND=Path(__file__).resolve().parents[1]
CACHE=Path('/mnt/data/research/bioinformatics-discovery/cache/jores2021')

def main():
    seq=pd.concat([pd.read_csv(CACHE/f'CNN__CNN_{s}_leaf.tsv',sep='\t',usecols=['gene','sequence']) for s in ['train','test']])
    seq=seq[seq.gene.str.match(r'^AT1G\d{5}$')].set_index('gene');assert seq.index.is_unique
    m=pd.read_pickle(CACHE/'014_full_estimators.pkl');m=m[(m.sp=='At')&(m.dna_cutoff==5)&m.gene.isin(seq.index)]
    valid=m.groupby('gene').agg(n=('rep','size'),minimum=('dna_barcodes','min'));valid=valid.index[(valid.n==4)&(valid.minimum>=10)]
    m=m[m.gene.isin(valid)].groupby(['gene','sys'])[['retained_median','all_pooled']].mean().unstack()
    m.columns=['_'.join(c) for c in m.columns];d=seq.join(m,how='inner').dropna();s=d.sequence
    f={b:s.str.count(b)/170 for b in 'CGT'}
    for motif in map(''.join,itertools.product('ACGT',repeat=2)):
        f[motif]=s.map(lambda x:sum(x[i:i+2]==motif for i in range(169))/169)
    gc=s.map(lambda x:(x.count('G')+x.count('C'))/170);f['GC2']=gc**2;f['GC3']=gc**3
    for lo,hi in [(0,60),(60,120),(120,165)]:f[f'GC_{lo}_{hi}']=s.map(lambda x:(x[lo:hi].count('G')+x[lo:hi].count('C'))/(hi-lo))
    f['TATA']=s.map(lambda x:bool(re.search(r'TATA[AT]A[AT][AG]',x[106:150]))).astype(int)
    for i in range(165,170):
        for b in 'CGT':f[f'base_{i}_{b}']=s.str[i].eq(b).astype(int)
    f['upstream_ATG_count']=s.str[135:165].map(lambda x:sum(x[i:i+3]=='ATG' for i in range(28)))
    f['terminal_ATG']=s.str[-5:].str.contains('ATG').astype(int)
    features=pd.DataFrame(f,index=d.index);features.to_csv(ROUND/'results/021_pilot_features.csv')
    summary={'promoters':len(d),'terminal_ATG_carriers':int(features.terminal_ATG.sum()),'effects':[],'advance':False}
    if summary['promoters']<500 or summary['terminal_ATG_carriers']<50:summary['decision']='Frozen carrier/sample gate failed; no effect estimate inspected.'
    else:
        x=np.column_stack([np.ones(len(d)),features.to_numpy()]);inv=np.linalg.pinv(x.T@x);lev=np.einsum('ij,jk,ik->i',x,inv,x)
        for host in ['leaf','proto']:
            for method in ['retained_median','all_pooled']:
                y=d[f'{method}_{host}'].to_numpy();beta=np.linalg.lstsq(x,y,rcond=None)[0];e=y-x@beta;w=(e/(1-lev))**2;cov=inv@(x.T@(x*w[:,None]))@inv
                se=float(np.sqrt(cov[-1,-1]));summary['effects'].append({'host':host,'method':method,'log2_effect':float(beta[-1]),'HC3_standard_error':se,'t':float(beta[-1]/se),'upstream_ATG_effect':float(beta[-2]),'design_rank':int(np.linalg.matrix_rank(x))})
        summary['advance']=all(a['log2_effect']<=-.5 and a['t']<=-3 for a in summary['effects'])
        summary['decision']='Advance to novelty review.' if summary['advance'] else 'Frozen effect gate failed; do not examine validation chromosomes or species.'
        d.drop(columns='sequence').to_csv(ROUND/'results/021_pilot_activities.csv')
    (ROUND/'results/021_summary.json').write_text(json.dumps(summary,indent=2)+'\n');print(json.dumps(summary,indent=2))

if __name__=='__main__':main()