#!/usr/bin/env python3
"""Add chr19 chromosome-map/target-profile embeddings and locus/link/ClinVar sets.

Inputs are existing chr19 CSVs (paths or URLs). Gene symbols resolve to the
matrix's complete SYMBOL|ENSEMBL|chr:TSS identifiers. Missing genes are reported,
never added as synthetic zero columns. Existing outputs require --overwrite.
"""
import argparse,csv,io,json,sys,urllib.request,hashlib,importlib.metadata
from pathlib import Path
import numpy as np
csv.field_size_limit(100_000_000)
COLUMNS=[('LDLR','LDLR',(11050000,11150000),(42,120,214)),('APOE–APOC','APOE',(44870000,44970000),(235,104,52)),('TM6SF2','TM6SF2',(19220000,19330000),(27,175,122)),('CYP2A6–2B6','CYP2A6',(40820000,41030000),(232,123,164))]
def read_csv(src):
    if src.startswith(('http://','https://')):
        with urllib.request.urlopen(src,timeout=120) as f:text=f.read().decode('utf-8-sig')
    else:text=Path(src).read_text(encoding='utf-8-sig')
    return list(csv.reader(io.StringIO(text))),hashlib.sha256(text.encode()).hexdigest()
def span(eid):
    chrom,pos=eid.rsplit(':',1);assert chrom=='chr19',eid
    start,end=map(int,pos.split('-'));assert 0<=start<end
    return start,end
def write_xy(p,ids,e):
    assert e.shape==(len(ids),2) and np.isfinite(e).all()
    with p.open('w',newline='') as f:
        w=csv.writer(f);w.writerow(['element_id','x','y'])
        w.writerows([i,format(float(x),'.9g'),format(float(y),'.9g')] for i,(x,y) in zip(ids,e))
def main():
    p=argparse.ArgumentParser(description=__doc__);p.add_argument('--matrix',default='chr19.matrix.csv');p.add_argument('--clinvar',default='chr19.clinvar.metadata.csv');p.add_argument('--out',default='.');p.add_argument('--overwrite',action='store_true');a=p.parse_args()
    out=Path(a.out);out.mkdir(parents=True,exist_ok=True)
    names=['chr19.map.embedding.csv','chr19.links.embedding.csv','chr19.sets.obsSets.json','chr19.extras.provenance.json']
    if not a.overwrite:
        for name in names:
            if (out/name).exists():raise FileExistsError(out/name)
    rows,mhash=read_csv(a.matrix);genes=rows[0][1:];ids=[r[0] for r in rows[1:]]
    assert len(ids)==len(set(ids)) and len(genes)==len(set(genes))
    X=np.array([r[1:] for r in rows[1:]],dtype=np.float32);assert X.shape==(len(ids),len(genes)) and np.isfinite(X).all() and (X>=0).all()
    cv,chash=read_csv(a.clinvar);ci=cv[0].index('element_id');ki=cv[0].index('clinvar_count');count={r[ci]:int(r[ki]) for r in cv[1:]}
    assert len(cv)-1==len(count) and set(count)==set(ids) and all(k>=0 for k in count.values())
    n=np.array([count[i] for i in ids]);spans=np.array([span(i) for i in ids]);mid=spans.mean(axis=1)/1e6
    y=np.log2(1+n);yscale=(mid.max()-mid.min())*.3/max(float(y.max()),1)
    norms=np.linalg.norm(X,axis=1,keepdims=True);assert (norms>0).all();Z=X/norms
    try:
        import umap
    except ImportError:
        U,S,_=np.linalg.svd(Z-Z.mean(axis=0),full_matrices=False);E=U[:,:2]*S[:2];method='PCA';params={'input':'L2-normalized gene-score profiles','components':2}
    else:
        params={'n_neighbors':25,'min_dist':.25,'metric':'cosine','random_state':0,'n_jobs':1,'input':'L2-normalized gene-score profiles'}
        E=umap.UMAP(**{k:v for k,v in params.items() if k!='input'}).fit_transform(Z);method='UMAP'
    E=E-E.mean(axis=0);scale=max(float(np.sqrt(np.mean(E**2))),1e-9);E=E/scale*10
    def inside(lo,hi):return [[eid,None] for eid,(s,e) in zip(ids,spans) if e>lo and s<hi]
    loci=[];hoods=[];linked=[];mapping={};missing=[]
    for name,g,(lo,hi),color in COLUMNS:
        loci.append({'name':name+' locus','color':list(color),'set':inside(lo,hi)})
        center=(lo+hi)//2;hoods.append({'name':name+' ±1 Mb','color':[int(v*.6+255*.4) for v in color],'set':inside(center-1000000,center+1000000)})
        js=[j for j,label in enumerate(genes) if label.split('|')[0]==g];mapping[g]=[genes[j] for j in js]
        if js:linked.append({'name':'linked to '+g,'color':list(color),'set':[[eid,None] for eid,v in zip(ids,(X[:,js]>0).any(axis=1)) if v]})
        else:missing.append(g)
    clinvar=[{'name':label,'color':color,'set':[[eid,None] for eid,k in zip(ids,n) if bool(k)==present]} for label,color,present in [('With ClinVar overlap',[255,140,0],True),('No ClinVar overlap',[110,110,110],False)]]
    tree=[{'name':'Gosling columns','children':loci+hoods},{'name':'rE2G links (score > 0)','children':linked},{'name':'ClinVar overlap','children':clinvar}]
    write_xy(out/names[0],ids,np.column_stack([mid,y*yscale]));write_xy(out/names[1],ids,E)
    (out/names[2]).write_text(json.dumps({'version':'0.1.3','datatype':'obs','tree':tree},indent=2,ensure_ascii=False)+'\n')
    report={'elements':len(ids),'genes':len(genes),'matrix_input':a.matrix,'matrix_text_sha256':mhash,'clinvar_input':a.clinvar,'clinvar_text_sha256':chash,'clinvar_positive_elements':int((n>0).sum()),'map':{'x':'element midpoint in Mb','y':'log2(1 + ClinVar count) times scale','y_scale':yscale,'raw_count_range':[int(n.min()),int(n.max())]},'landscape':{'method':method,'parameters':params,'postprocessing':'center both axes and apply one common scale factor; geometry preserved','packages':{p:importlib.metadata.version(p) for p in ['numpy','scikit-learn']+(['umap-learn'] if method=='UMAP' else [])}},'gene_feature_ids':mapping,'missing_requested_genes':missing,'locus_windows':{name:[lo,hi] for name,g,(lo,hi),col in COLUMNS},'set_sizes':{group['name']:{node['name']:len(node['set']) for node in group['children']} for group in tree},'interpretation':'Exploratory embedding of predicted target-link profiles, not expression, physical distance, or evidence of causal regulation.'}
    (out/names[3]).write_text(json.dumps(report,indent=2,ensure_ascii=False)+'\n');print(json.dumps(report,indent=2,ensure_ascii=False))
if __name__=='__main__':main()