#!/usr/bin/env python3
"""Attempt 016: fixed composition-adjusted motif contrast in maize-origin promoters."""
import itertools,json,re
from pathlib import Path
import numpy as np
import pandas as pd
from scipy.stats import t as student_t
ROUND=Path(__file__).resolve().parents[1];CACHE=Path('/mnt/data/research/bioinformatics-discovery/cache/jores2021')
G4=re.compile(r'G{3,}(?:[ACGT]{1,7}G{3,}){3}')
C4=re.compile(r'C{3,}(?:[ACGT]{1,7}C{3,}){3}')

def main():
    sequences=pd.concat([pd.read_csv(CACHE/f'CNN__CNN_{split}_leaf.tsv',sep='\t',usecols=['gene','sp','sequence']) for split in ['train','test']])
    sequences=sequences[sequences.sp=='Zm'];assert not sequences.gene.duplicated().any()
    measurements=pd.read_pickle(CACHE/'014_full_estimators.pkl');measurements=measurements[(measurements.sp=='Zm')&(measurements.dna_cutoff==5)]
    valid=measurements.groupby('gene').agg(n=('rep','size'),minimum_barcodes=('dna_barcodes','min'))
    valid=valid.index[(valid.n==4)&(valid.minimum_barcodes>=10)]
    values=measurements[measurements.gene.isin(valid)].groupby(['gene','sys'])[['retained_median','all_pooled']].mean().unstack()
    values.columns=values.columns.to_flat_index()
    data=sequences.set_index('gene').join(values,how='inner').dropna();seq=data.sequence
    features={}
    for base in 'CGT':features[base]=seq.str.count(base)/170
    for pair in map(''.join,itertools.product('ACGT',repeat=2)):
        features[pair]=seq.map(lambda s:sum(s[i:i+2]==pair for i in range(len(s)-1))/(len(s)-1))
    gc=seq.map(lambda s:(s.count('G')+s.count('C'))/len(s))
    features['GC2']=gc**2;features['GC3']=gc**3
    for lo,hi in [(0,60),(60,120),(120,170)]:features[f'GC_{lo}_{hi}']=seq.map(lambda s:(s[lo:hi].count('G')+s[lo:hi].count('C'))/(hi-lo))
    features['TATA']=seq.map(lambda s:bool(re.search(r'TATA[AT]A[AT][AG]',s[106:150]))).astype(int)
    features['G4']=seq.map(lambda s:bool(G4.search(s))).astype(int);features['C4']=seq.map(lambda s:bool(C4.search(s))).astype(int)
    frame=pd.DataFrame(features);counts=frame[['G4','C4']].sum().to_dict();summary=dict(promoters=len(data),motif_counts={k:int(v) for k,v in counts.items()},results=[])
    if min(counts.values())<50:
        summary['advance']=False;summary['decision']='Insufficient canonical motif carriers for the predeclared pilot.'
    else:
        x=np.column_stack([np.ones(len(frame)),frame.to_numpy()]);rank=np.linalg.matrix_rank(x);inverse=np.linalg.pinv(x.T@x);leverage=np.einsum('ij,jk,ik->i',x,inverse,x)
        contrast=np.zeros(x.shape[1]);contrast[-2]=1;contrast[-1]=-1
        for method in ['retained_median','all_pooled']:
            y=(data[(method,'leaf')]-data[(method,'proto')]).to_numpy();beta=np.linalg.lstsq(x,y,rcond=None)[0];resid=y-x@beta
            weights=(resid/(1-leverage))**2;cov=inverse@(x.T@(x*weights[:,None]))@inverse
            effect=float(contrast@beta);se=float(np.sqrt(contrast@cov@contrast));t=effect/se
            summary['results'].append(dict(method=method,contrast=effect,standard_error_HC3=se,t=t,p_two_sided=float(2*student_t.sf(abs(t),len(y)-rank)),design_rank=int(rank)))
        results=summary['results'];summary['advance']=all(abs(r['contrast'])>=.25 and abs(r['t'])>=3 for r in results) and results[0]['contrast']*results[1]['contrast']>0
    (ROUND/'results/016_g4_summary.json').write_text(json.dumps(summary,indent=2)+'\n')
    frame.to_csv(ROUND/'results/016_maize_sequence_features.csv')
    print(json.dumps(summary,indent=2))

if __name__=='__main__':main()