#!/usr/bin/env python3
"""Attempt 017: frozen chromosome-1 promoter/terminator coordination pilot."""
import itertools,json
from pathlib import Path
import numpy as np
import pandas as pd
from scipy.stats import spearmanr,pearsonr,rankdata
ROUND=Path(__file__).resolve().parents[1]
CACHE=Path('/mnt/data/research/bioinformatics-discovery/cache')

def composition(seq):
    f={b:seq.str.count(b)/170 for b in 'CGT'}
    for k in map(''.join,itertools.product('ACGT',repeat=2)):
        f[k]=seq.map(lambda s:sum(s[i:i+2]==k for i in range(169))/169)
    gc=seq.map(lambda s:(s.count('G')+s.count('C'))/170)
    f['GC2']=gc**2;f['GC3']=gc**3
    for lo,hi in [(0,60),(60,120),(120,170)]:
        f[f'GC_{lo}_{hi}']=seq.map(lambda s:(s[lo:hi].count('G')+s[lo:hi].count('C'))/(hi-lo))
    return pd.DataFrame(f)

def main():
    prom=pd.concat([pd.read_csv(CACHE/'jores2021'/f'CNN__CNN_{s}_leaf.tsv',sep='\t',usecols=['gene','sequence']) for s in ['train','test']])
    prom=prom[prom.gene.str.match(r'^AT1G\d{5}$')].set_index('gene')
    m=pd.read_pickle(CACHE/'jores2021/014_full_estimators.pkl');m=m[(m.sp=='At')&(m.dna_cutoff==5)&m.gene.isin(prom.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=['prom_'+a+'_'+b for a,b in m.columns]
    term=pd.read_csv(CACHE/'round-002/terminators2024/CNN__terminator_data.tsv',sep='\t')
    term=term[term.id.str.match(r'^AT1G\d{5}(?:_[0-9]+)?$')].copy();term['gene']=term.id.str.extract(r'^(AT1G\d{5})')
    # Average features across native isoform terminators, rather than choosing by activity.
    features=composition(term.sequence);features['gene']=term.gene;features=features.groupby('gene').mean().add_prefix('term_')
    term=term.groupby('gene')[['enrichment_tobacco','enrichment_maize']].mean().rename(columns={'enrichment_tobacco':'term_leaf','enrichment_maize':'term_proto'})
    data=prom.join(m,how='inner').join(term,how='inner').dropna()
    cov=composition(data.sequence).add_prefix('prom_').join(features,how='inner').loc[data.index]
    x=np.column_stack([np.ones(len(cov)),np.apply_along_axis(rankdata,0,cov.to_numpy())]);rank=int(np.linalg.matrix_rank(x))
    rows=[]
    for host in ['leaf','proto']:
        for method in ['retained_median','all_pooled']:
            a=data[f'prom_{method}_{host}'].to_numpy();b=data[f'term_{host}'].to_numpy()
            ra=rankdata(a);rb=rankdata(b);ea=ra-x@np.linalg.lstsq(x,ra,rcond=None)[0];eb=rb-x@np.linalg.lstsq(x,rb,rcond=None)[0]
            rows.append(dict(host=host,method=method,n=len(data),spearman=float(spearmanr(a,b).statistic),partial_rank_r=float(pearsonr(ea,eb).statistic)))
    summary=dict(design_rank=rank,results=rows,advance=bool(len(data)>=500 and all(r['partial_rank_r']>=.15 for r in rows)))
    data.drop(columns='sequence').to_csv(ROUND/'results/017_chr1_matched_activities.csv')
    (ROUND/'results/017_summary.json').write_text(json.dumps(summary,indent=2)+'\n');print(json.dumps(summary,indent=2))

if __name__=='__main__':main()