#!/usr/bin/env python3
"""Build chr4 gene, cCRE, and ClinVar files from HepG2 ENCFF960GTR."""
import copy
import csv
import gzip
import hashlib
import json
import subprocess
from update_clinvar import update_clinvar
from pathlib import Path
from urllib.parse import unquote

ROOT = Path(__file__).resolve().parent

def write_csv(name, fields, rows):
    with (ROOT / name).open('w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)

def digest(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()

BASE_CONFIG = {'version': '1.0.18', 'name': 'HepG2 | chr4 | cCRE overlap + gene scores', 'description': 'Chr4 only. Left: binary overlap of each regulatory element with eight cCRE classes (1 = overlap, 0 = no overlap). Right: original element-by-gene rE2G scores and gene list. Rows retain the source chr4 matrix order; row navigation and highlighting are shared. cCRE annotations describe interval overlaps, not cCRE-specific gene predictions. Zero gene scores mean no reported link in the thresholded input.', 'initStrategy': 'auto', 'datasets': [{'uid': 'HepG2', 'name': 'HepG2 OG chr4', 'files': [{'fileType': 'obsFeatureMatrix.csv', 'url': 'https://users.abdenlab.org/voy/vitessce_cvh/hepg2_sci/chr4.matrix.csv', 'coordinationValues': {'obsType': 'element', 'featureType': 'gene', 'featureValueType': 'score'}}, {'fileType': 'obsLabels.csv', 'url': 'https://users.abdenlab.org/voy/vitessce_cvh/hepg2_sci/chr4.ccre_metadata.csv', 'coordinationValues': {'obsType': 'element', 'obsLabelsType': 'Overlapping cCREs (ID: class)'}, 'options': {'obsIndex': 'element_id', 'obsLabels': 'ccre_id_class'}}, {'fileType': 'obsLabels.csv', 'url': 'https://users.abdenlab.org/voy/vitessce_cvh/hepg2_sci/chr4.ccre_metadata.csv', 'coordinationValues': {'obsType': 'element', 'obsLabelsType': 'Overlapping cCRE count'}, 'options': {'obsIndex': 'element_id', 'obsLabels': 'ccre_count'}}, {'fileType': 'obsFeatureMatrix.csv', 'url': 'https://users.abdenlab.org/voy/vitessce_cvh/hepg2_sci/chr4.ccre.matrix.csv', 'coordinationValues': {'obsType': 'element', 'featureType': 'cCRE class', 'featureValueType': 'overlap'}}]}], 'coordinationSpace': {'dataset': {'A': 'HepG2'}, 'obsType': {'A': 'element'}, 'featureType': {'A': 'gene', 'classes': 'cCRE class'}, 'featureValueType': {'A': 'score', 'classes': 'overlap'}, 'featureSelection': {'A': None, 'classes': None}, 'featureHighlight': {'A': None, 'classes': None}, 'obsHighlight': {'A': None}, 'obsLabelsType': {'ccre_matches': 'Overlapping cCREs (ID: class)', 'ccre_count': 'Overlapping cCRE count'}, 'heatmapZoomX': {'rows': 0}, 'heatmapZoomY': {'rows': 0}, 'heatmapTargetY': {'rows': 0}, 'heatmapTargetX': {'scores': 0, 'classes': 0}, 'featureValueColormap': {'scores': 'plasma', 'classes': 'viridis'}, 'featureValueColormapRange': {'scores': [0, 1], 'classes': [0, 1]}}, 'layout': [{'component': 'heatmap', 'coordinationScopes': {'dataset': 'A', 'obsType': 'A', 'featureType': 'classes', 'featureValueType': 'classes', 'featureSelection': 'classes', 'featureHighlight': 'classes', 'obsHighlight': 'A', 'obsLabelsType': ['ccre_matches', 'ccre_count'], 'heatmapZoomX': 'rows', 'heatmapZoomY': 'rows', 'heatmapTargetY': 'rows', 'heatmapTargetX': 'classes', 'featureValueColormap': 'classes', 'featureValueColormapRange': 'classes'}, 'props': {'transpose': False, 'title': 'HepG2 chr4: cCRE class overlap (0/1)', 'variablesLabelOverride': 'cCRE class', 'observationsLabelOverride': 'Regulatory element'}, 'x': 0, 'y': 0, 'w': 4, 'h': 12}, {'component': 'heatmap', 'coordinationScopes': {'dataset': 'A', 'obsType': 'A', 'featureType': 'A', 'featureValueType': 'A', 'featureSelection': 'A', 'featureHighlight': 'A', 'obsHighlight': 'A', 'obsLabelsType': ['ccre_matches', 'ccre_count'], 'heatmapZoomX': 'rows', 'heatmapZoomY': 'rows', 'heatmapTargetY': 'rows', 'heatmapTargetX': 'scores', 'featureValueColormap': 'scores', 'featureValueColormapRange': 'scores'}, 'props': {'transpose': False, 'title': 'HepG2 chr4: element × gene', 'observationsLabelOverride': 'Regulatory element'}, 'x': 4, 'y': 0, 'w': 6, 'h': 12}, {'component': 'featureList', 'coordinationScopes': {'dataset': 'A', 'obsType': 'A', 'featureType': 'A', 'featureValueType': 'A', 'featureSelection': 'A', 'featureHighlight': 'A', 'obsHighlight': 'A'}, 'x': 10, 'y': 0, 'w': 2, 'h': 12}]}

SOURCE = Path('/net/users/voy/nmf/data_itcr/cnTFeffects/reanalysis0612tecounts/external/encode_rE2G_full/beds/ENCFF960GTR.bed.gz')
CCRE_REFERENCE = Path('/net/users/voy/nmf/data_itcr/configs/GRCh38-cCREs.bed')


def build_inputs():
    with gzip.open(SOURCE, 'rt') as f:
        links = [r for r in csv.DictReader(f, delimiter='\t') if r['#chr'] == 'chr4']
    assert all(r['CellType'] == 'HepG2' for r in links)
    elements = sorted({(int(r['start']), int(r['end'])) for r in links})
    genes = sorted({(int(r['TargetGeneTSS']), r['TargetGeneEnsemblID'], r['TargetGene']) for r in links})
    values = {}
    for r in links:
        e = (int(r['start']), int(r['end']))
        g = (int(r['TargetGeneTSS']), r['TargetGeneEnsemblID'], r['TargetGene'])
        assert (e, g) not in values, 'Duplicate element/gene/TSS link'
        assert 0 < float(r['Score']) <= 1
        values[e, g] = r['Score']
    labels = [f'{symbol}|{ens}|chr4:{tss}' for tss, ens, symbol in genes]
    assert {'ALB', 'AFP'} <= {g[2] for g in genes}
    ids = [f'chr4:{a}-{b}' for a, b in elements]
    with (ROOT / 'chr4.matrix.csv').open('w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(['element_id', *labels])
        writer.writerows([eid, *[values.get((e, g), '0') for g in genes]] for eid, e in zip(ids, elements))
    # Verify every nonzero score and every zero against the input links.
    with (ROOT / 'chr4.matrix.csv').open() as f:
        reader = csv.reader(f)
        assert next(reader) == ['element_id', *labels]
        for e, eid in zip(elements, ids):
            row = next(reader)
            assert row[0] == eid
            assert row[1:] == [values.get((e, g), '0') for g in genes]
        assert next(reader, None) is None
    bed = ''.join(f'chr4\t{a}\t{b}\t{eid}\n' for (a, b), eid in zip(elements, ids))
    result = subprocess.run(['bedtools', 'intersect', '-a', 'stdin', '-b', str(CCRE_REFERENCE), '-wao'],
                            input=bed, text=True, capture_output=True, check=True)
    overlaps = {eid: [] for eid in ids}
    for line in result.stdout.splitlines():
        f = line.split('\t')
        assert len(f) == 11
        if int(f[-1]) > 0:
            overlaps[f[3]].append((int(f[5]), int(f[6]), f[7], f[8], f[9], int(f[10])))
    classes = ['PLS', 'pELS', 'dELS', 'CA-H3K4me3', 'CA-CTCF', 'CA-TF', 'CA', 'TF']
    metadata = []
    for (a, b), eid in zip(elements, ids):
        matches = sorted(overlaps[eid])
        assert all(m[4] in classes for m in matches)
        metadata.append(dict(element_id=eid, chrom='chr4', start=a, end=b, ccre_count=len(matches),
            ccre_ids=';'.join(m[3] for m in matches), ccre_classes=';'.join(m[4] for m in matches),
            ccre_id_class='; '.join(f'{m[3]}: {m[4]}' for m in matches) or 'No overlapping cCRE',
            ccre_intervals=';'.join(f'chr4:{m[0]}-{m[1]}' for m in matches),
            ccre_overlap_bp=';'.join(str(m[5]) for m in matches), ccre_source_ids=';'.join(m[2] for m in matches)))
    write_csv('chr4.ccre_metadata.csv', list(metadata[0]), metadata)
    write_csv('chr4.ccre.matrix.csv', ['element_id', *classes],
              [dict(element_id=eid, **{c: int(any(m[4] == c for m in overlaps[eid])) for c in classes}) for eid in ids])
    config = copy.deepcopy(BASE_CONFIG)
    config['name'] = 'HepG2 ENCFF960GTR | chr4 | cCRE + gene scores'
    config['datasets'][0]['name'] = 'HepG2 ENCFF960GTR chr4'
    config['description'] = ('HepG2 ENCODE rE2G ENCFF960GTR / ENCSR411RGL, GRCh38 chr4 only. '
        'Original Score values are preserved; zero denotes an unreported source link. '
        'Elements and gene/TSS columns follow genomic order. cCRE classes represent positive-length '
        'overlaps with the GRCh38 registry, not cCRE-specific gene predictions. '
        'ALB and AFP are included as reported target genes. This is regulatory-link data, not expression data.')
    report = dict(source_accession='ENCFF960GTR', source_dataset='ENCSR411RGL', source_file=str(SOURCE),
        source_sha256=digest(SOURCE), chromosome='chr4', elements=len(elements), gene_columns=len(genes),
        source_links=len(links), all_scores_verified_against_source=True, ccre_classes=len(classes),
        ccre_reference=str(CCRE_REFERENCE), ccre_reference_sha256=digest(CCRE_REFERENCE),
        marker_gene_columns=[label for label in labels if label.split('|')[0] in ('ALB', 'AFP')],
        marker_link_counts={s: sum(r['TargetGene'] == s for r in links) for s in ['ALB', 'AFP']},
        browser_rendering_verified=False)
    (ROOT / 'validation.json').write_text(json.dumps(report, indent=2) + '\n')
    return config


def main():
    config = build_inputs()
    report = update_clinvar(ROOT, 'chr4')
    children = json.loads((ROOT / 'chr4.clinvar.obsSets.json').read_text())['tree'][0]['children']
    files = config['datasets'][0]['files']
    files.append(dict(fileType='obsFeatureMatrix.csv', url='chr4.clinvar.matrix.csv',
        coordinationValues=dict(obsType='element', featureType='ClinVar annotation', featureValueType='variant count')))
    files.append(dict(fileType='obsSets.json', url='chr4.clinvar.obsSets.json', coordinationValues=dict(obsType='element')))
    cs = config['coordinationSpace']
    for scope, label, column in [('clinvar_count', 'ClinVar overlap count', 'clinvar_count'),
                                 ('clinvar_details', 'Overlapping ClinVar variants', 'clinvar_details')]:
        files.append(dict(fileType='obsLabels.csv', url='chr4.clinvar.metadata.csv',
            coordinationValues=dict(obsType='element', obsLabelsType=label),
            options=dict(obsIndex='element_id', obsLabels=column)))
        cs['obsLabelsType'][scope] = label
    for key, value in dict(featureType='ClinVar annotation', featureValueType='variant count',
            featureSelection=None, featureHighlight=None, heatmapTargetX=0,
            featureValueColormap='viridis', featureValueColormapRange=[0, 1]).items():
        cs[key]['clinvar'] = value
    cs['obsSetSelection'] = {'clinvar': []}
    cs['obsSetColor'] = {'clinvar': [dict(path=['ClinVar overlap', c['name']], color=c['color']) for c in children]}
    ccre, gene, gene_list = config['layout']
    strip = copy.deepcopy(ccre)
    for key in ['featureType', 'featureValueType', 'featureSelection', 'featureHighlight',
                'heatmapTargetX', 'featureValueColormap', 'featureValueColormapRange']:
        strip['coordinationScopes'][key] = 'clinvar'
    strip['props'].update(title='ClinVar overlap count', variablesLabelOverride='ClinVar annotation')
    ccre.update(x=0, w=3)
    strip.update(x=3, w=2)
    gene.update(x=5, w=5)
    gene_list.update(x=10, w=2, h=8)
    for view in [ccre, strip, gene]:
        # Keep lengthy variant descriptions in the metadata CSV, not hover popups.
        view['coordinationScopes']['obsLabelsType'] = {
            'classes': ['ccre_matches', 'ccre_count'],
            'clinvar': ['clinvar_count'],
            'A': ['ccre_count', 'clinvar_count'],
        }[view['coordinationScopes']['featureType']]
        view['coordinationScopes'].update(obsSetSelection='clinvar', obsSetColor='clinvar')
    group_view = dict(component='obsSets', coordinationScopes=dict(dataset='A', obsType='A',
        obsSetSelection='clinvar', obsSetColor='clinvar'), x=10, y=8, w=2, h=4)
    # NBSP is truthy and remains blank after Vitessce's pluralization.
    strip['props']['observationsLabelOverride'] = '\u00a0'
    gene['props']['observationsLabelOverride'] = '\u00a0'
    cs['featureValueColormap'].update(classes='jet', clinvar='jet')
    config['layout'] = [ccre, strip, gene, gene_list, group_view]
    config['name'] = 'HepG2 ENCFF960GTR | chr4 | cCRE + ClinVar overlap + gene scores'
    config['description'] += (' ClinVar source: clinvar_hg38_all.vcf.gz (all clinical-significance categories, with no phenotype filter). ClinVar strip: number of supplied chr4 ClinVar records overlapping each element. '
        'Overlap uses [POS-1, POS-1+len(REF)), including indel anchor bases. Counts are stored as raw integers; '
        'hover labels provide exact counts. ClinVar associations do not establish HepG2 genotypes or regulatory effects. '
        'VCF annotated genes are not substituted for rE2G target genes. All rows are shown initially; overlap groups are selectable.')
    for dataset in config['datasets']:
        for entry in dataset['files']:
            entry['url'] = 'https://users.abdenlab.org/voy/vitessce_cvh/hepg2_sci/' + entry['url'].rsplit('/', 1)[-1]
    (ROOT / 'chr4.json').write_text(json.dumps(config, indent=2) + '\n')
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()