#!/usr/bin/env python3
"""Candidate 012. RNA-only thinning; fixed observed DNA counts and mapping."""
from coding_barcode_pilot import ROOT,CACHE,read_counts,np,pd
import json

def main():
    sub=pd.read_csv(CACHE/'analysis__subassembly__subassembly_At.tsv',sep='\t',low_memory=False)
    sub=sub[(sub.variant=='WT')&sub.FL&(sub.gene!='35Spr')]
    prefix='data__barcode_counts__leaf__At_Rep1__barcodes_pPSup_AtPRO_35SEnh_I'
    dna=read_counts(CACHE/f'{prefix}_inp.count.gz','dna')
    rna=read_counts(CACHE/f'{prefix}_dark.count.gz','rna')
    data=dna[dna.dna>=5].merge(sub,on='barcode',validate='one_to_one').merge(rna,on='barcode',how='left',validate='one_to_one')
    data['rna']=data.rna.fillna(0).astype(int)
    support=data.groupby('gene').barcode.size();eligible=support.index[support>=10]
    data=data[data.gene.isin(eligible)].copy()
    rng=np.random.default_rng(20260919)
    estimates=[]
    for fraction in [1.,.5,.25]:
        counts=data.rna.to_numpy() if fraction==1 else rng.binomial(data.rna.to_numpy(),fraction)
        temp=data.assign(observed_rna=counts)
        selected=temp[temp.observed_rna>=5].copy()
        selected['ratio']=np.log2(selected.observed_rna/selected.dna)-np.log2(fraction)
        original=selected.groupby('gene').ratio.median().rename('retained_median')
        pooled=temp.groupby('gene')[['observed_rna','dna']].sum()
        # 0.5 prevents log(0); with >=10 DNA-supported barcodes this is a small
        # aggregate-level pseudocount. It is applied at every depth identically.
        pooled['pooled']=np.log2((pooled.observed_rna+.5)/(pooled.dna+.5))-np.log2(fraction)
        df=pooled[['pooled']].join(original);df['fraction']=fraction
        df['surviving_barcodes']=selected.groupby('gene').barcode.size()
        df['dna_barcodes']=support;estimates.append(df.reset_index())
    all_estimates=pd.concat(estimates,ignore_index=True)
    all_estimates.to_csv(ROOT/'results/012_depth_estimates.csv',index=False)
    rows=[]
    for method in ['retained_median','pooled']:
        wide=all_estimates.pivot(index='gene',columns='fraction',values=method)
        for fraction in [.5,.25]:
            pair=wide[[1.,fraction]].dropna();delta=pair[fraction]-pair[1.]
            rows.append(dict(method=method,fraction=fraction,full_depth_promoters=int(wide[1.].notna().sum()),
                paired_promoters=len(pair),spearman=float(pair.corr(method='spearman').iloc[0,1]),
                median_shift=float(delta.median()),median_absolute_shift=float(delta.abs().median()),
                fraction_absolute_shift_over_1=float((delta.abs()>1).mean())))
    (ROOT/'results/012_depth_summary.json').write_text(json.dumps(rows,indent=2)+'\n')
    print(json.dumps(rows,indent=2))

if __name__=='__main__':main()