#!/usr/bin/env python3
"""Exact and reverse-complement overlap audit; no model fitting."""
from collections import Counter,defaultdict
import hashlib
import json
from pathlib import Path
import pandas as pd

ROOT=Path(__file__).resolve().parents[1]
CACHE=Path('/mnt/data/research/bioinformatics-discovery/cache')
COMP=str.maketrans('ACGTN','TGCAN')
def canon(s):
    s=s.upper();return min(s,s.translate(COMP)[::-1])

def audit(frames,name):
    out={'dataset':name,'splits':{},'overlap':{}}
    sets={}
    for split,df in frames.items():
        df=df.copy();df['canonical']=df.sequence.map(canon);frames[split]=df
        sets[split]=set(df.canonical)
        bylabel=df.groupby('canonical').label.nunique()
        out['splits'][split]={'rows':len(df),'unique_sequences':df.sequence.nunique(),
            'unique_canonical':df.canonical.nunique(),'sequence_lengths':df.sequence.str.len().value_counts().to_dict(),
            'canonical_label_conflicts':int((bylabel>1).sum()),'labels':df.label.value_counts().head(10).to_dict()}
    for a,b in [('train','test'),('train','validation'),('validation','test')]:
        common=sets[a]&sets[b]
        hits=frames[b][frames[b].canonical.isin(common)]
        out['overlap'][a+'__'+b]={'unique_canonical_shared':len(common),'affected_rows_in_second':len(hits),
            'fraction_of_second':len(hits)/len(frames[b])}
        if len(hits):
            # Save identifiers, labels and sequence hashes, not the full sequences.
            examples=hits.drop(columns=['sequence','canonical']).copy()
            examples['sequence_sha256']=hits.sequence.map(lambda s:hashlib.sha256(s.encode()).hexdigest())
            examples.to_csv(ROOT/f'results/007_{name}_{a}_{b}_overlap.csv',index=False)
    return out

def fasta(path):
    out=[];header=None;parts=[]
    for line in path.read_text().splitlines():
        if line.startswith('>'):
            if header is not None:out.append((header,''.join(parts)))
            header=line[1:];parts=[]
        else:parts.append(line.strip())
    if header is not None:out.append((header,''.join(parts)))
    return out

def main():
    frames={s:pd.read_csv(CACHE/'plant_promoters'/f'{filename}.csv') for s,filename in [('train','train'),('test','test'),('validation','dev')]}
    results=[audit(frames,'pdllm_core')]
    print(json.dumps(results,indent=2))
    for task in ['promoter_strength_leaf']:
        files={s:CACHE/'pgb'/f'{task}_{s}.fa' for s in ['train','test','validation']}
        if not all(f.exists() for f in files.values()):continue
        frames={}
        for split,file in files.items():
            records=fasta(file);print('Header examples:',[r[0] for r in records[:3]])
            frames[split]=pd.DataFrame([dict(identifier=h.rsplit("|",1)[0],sequence=s,label=float(h.rsplit("|",1)[1])) for h,s in records])
        results.append(audit(frames,task))
    (ROOT/'results/007_promoter_split_audit.json').write_text(json.dumps(results,indent=2))
    print(json.dumps(results,indent=2))

if __name__=='__main__':main()