#!/usr/bin/env python3
"""Generate the trackDb .ra files for the transcription start site tracks of one
assembly, one file per top-level track:

    transcriptionStart.ra   the superTrack both data sources sit in
    encode4ProCap.ra        the experimental PRO-cap composite, hg38 only
    proCapNet.ra            the model predictions and, on hg38, the
                            sequence-contribution scores

Each is included from the assembly's trackDb.ra with an alpha release tag, so
the whole set is on alpha and nothing reaches beta or public.

These are traditional composites, not faceted ones.  A faceted composite is
routed to facetedCompositeUi(), which returns before cfgByCfgType(), so the
wiggle controls are never drawn; signal tracks need them.  A container multiWig
under a traditional composite still draws as one overlay row per cell line.
What it does not get is an inline config block on the composite page, and
deselecting both strands of a cell line leaves an empty row where the overlay
was, both because compositeUiSubtracks() lists descendant leaves and a multiWig
is not a leaf.  That is #38441, to be fixed in hgTrackUi.

An earlier attempt made each data source a superTrack holding the overlays
directly, which gave every overlay its own configuration page and sidestepped
#38441.  It was reverted: it cost the subtrack matrix and the sample class
filter, and nesting the two sources in transcriptionStart then needed the
superTrack-in-superTrack fix of #38460.

Subtrack names are <composite>_<cell>_<dataType>.
"""
import argparse
import os.path as osp
import re
from pycbio.sys import cli, fileOps
from pycbio.tsv import TsvReader

GBDB = "/gbdb/{db}/{track}"


def parseArgs():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("db")
    parser.add_argument("experimentsTsv")
    parser.add_argument("outDir",
                        help="trackDb directory for db, one .ra written per top-level track")
    return cli.parseOptsArgsWithLogging(parser)

def subGroupTag(name):
    """A subGroup tag has to be a bare identifier, but cell line names carry
    punctuation (Caco-2).  Only the tag is sanitized; the label keeps the real
    name."""
    return re.sub(r"[^A-Za-z0-9]", "", name)

def stanza(indent, lines):
    pad = " " * indent
    return "".join(pad + line + "\n" for line in lines) + "\n"

def writeRa(db, outDir, track, contents):
    outRa = osp.join(outDir, track + ".ra")
    fileOps.ensureFileDir(outRa)
    with fileOps.AtomicFileOpen(outRa) as fh:
        print(f"# Generated by hg/makeDb/outside/proCapNet/proCapNetTrackDb for {db}.",
              file=fh)
        print("# Do not edit by hand, edit the script and regenerate.\n", file=fh)
        fh.write(contents)

def superStanza():
    return stanza(0, [
        "track transcriptionStart",
        "superTrack on",
        "group rna",
        "shortLabel Transcription Initiation (TSS)",
        "longLabel Transcription initiation (TSS)",
    ])

def compositeStanza(track, shortLabel, longLabel, experiments, dataTypes, priority):
    """The matrix runs over the leaf bigWigs, which is the level hgTrackDb checks
    and the level compositeUiSubtracks() lists, so each strand is a cell of its
    own.  The multiWig containers carry no subGroups: a container spans both
    strands, so it has no single value for the data type dimension, and
    hgTrackDb -strict rejects a leaf whose groups the parent does not declare.

    dataTypes is a list of (tag, label) pairs, the matrix's X dimension."""
    cells = " ".join(f"{subGroupTag(e.cell)}={e.cell}" for e in experiments)
    types = " ".join(f"{tag}={label}" for tag, label in dataTypes)
    classes = " ".join(f"{subGroupTag(c)}={c}" for c in
                       sorted({e.sampleClass for e in experiments}))
    return stanza(0, [
        f"track {track}",
        "parent transcriptionStart",
        "compositeTrack on",
        "type bigWig",
        f"shortLabel {shortLabel}",
        f"longLabel {longLabel}",
        f"subGroup1 cellLine Cell_line {cells}",
        f"subGroup2 dataType Data_type {types}",
        f"subGroup3 sampleClass Sample_class {classes}",
        "dimensions dimensionX=dataType dimensionY=cellLine dimA=sampleClass",
        "filterComposite dimA",
        "sortOrder cellLine=+ dataType=+",
        # a config wrench on every subtrack, so one cell line can be rescaled
        # without touching the rest
        "configurable on",
        "autoScale group",
        "alwaysZero on",
        "maxHeightPixels 100:40:8",
        "windowingFunction maximum",
        "noInherit on",
        "visibility hide",
        f"priority {priority}",
    ])

def strandOverlay(db, composite, exp, dataType, subDir, fileTag, shortLabel, longLabel,
                  priority):
    """One cell line's two strands as an overlay, plus strand drawn up and minus
    strand drawn down.  Both sources store the minus strand negated, so neither
    needs the trackDb negateValues setting.  That setting used to be on the
    ProCapNet minus tracks, and it broke the composite's own negate control: that
    control sets one shared value for the composite, which replaced the per-track
    setting and sent both strands the same way, with no route back to the default
    short of a cart reset.

    fileTag is the part of the file name that says what the file holds, so a
    bigWig downloaded on its own still names its source: the model for the
    predictions, the ENCODE experiment accession for the measurements.

    The parent line says "on", which makes these default-on.  The contribution
    scores stay off, so opening the composite gives signal rather than every data
    type at once."""
    gbdb = GBDB.format(db=db, track=composite) + subDir
    out = stanza(4, [
        f"track {composite}_{exp.cell}_{dataType}",
        f"parent {composite} on",
        "container multiWig",
        "aggregate solidOverlay",
        "showSubtrackColorOnUi on",
        "type bigWig",
        f"color {exp.color}",
        f"shortLabel {shortLabel}",
        f"longLabel {longLabel}",
        "onlyVisibility full",
        f"priority {priority}",
    ])
    out += stanza(8, [
        f"track {composite}_{exp.cell}_{dataType}_pos",
        f"parent {composite}_{exp.cell}_{dataType}",
        f"subGroups cellLine={subGroupTag(exp.cell)} dataType={dataType}Pos"
        f" sampleClass={subGroupTag(exp.sampleClass)}",
        "type bigWig",
        f"bigDataUrl {gbdb}/{exp.cell}.{fileTag}.pos.bw",
        f"color {exp.color}",
        f"shortLabel {shortLabel} +",
        f"longLabel {longLabel}, plus strand",
    ])
    out += stanza(8, [
        f"track {composite}_{exp.cell}_{dataType}_neg",
        f"parent {composite}_{exp.cell}_{dataType}",
        f"subGroups cellLine={subGroupTag(exp.cell)} dataType={dataType}Neg"
        f" sampleClass={subGroupTag(exp.sampleClass)}",
        "type bigWig",
        f"bigDataUrl {gbdb}/{exp.cell}.{fileTag}.neg.bw",
        f"color {exp.color}",
        f"altColor {exp.color}",
        f"shortLabel {shortLabel} -",
        f"longLabel {longLabel}, minus strand",
    ])
    return out

def contribStanza(db, exp, priority):
    gbdb = GBDB.format(db=db, track="proCapNet")
    return stanza(4, [
        f"track proCapNet_{exp.cell}_contrib",
        f"subGroups cellLine={subGroupTag(exp.cell)} dataType=contrib"
        f" sampleClass={subGroupTag(exp.sampleClass)}",
        "parent proCapNet off",
        "type bigWig",
        f"bigDataUrl {gbdb}/contrib/{exp.cell}.proCapNet-contrib.bw",
        "logo on",
        f"color {exp.color}",
        f"shortLabel {exp.cell} Contribution",
        f"longLabel {exp.cell} ({exp.modelAcc}) ProCapNet sequence-contribution scores",
        "onlyVisibility full",
        f"priority {priority}",
    ])

def proCapNetComposite(db, experiments):
    dataTypes = [("predPos", "Predicted_+"), ("predNeg", "Predicted_-")]
    if db == "hg38":
        dataTypes.append(("contrib", "Contribution_scores"))
    out = compositeStanza("proCapNet", "ProCapNet",
                          "ProCapNet predicted PRO-cap",
                          experiments, dataTypes, 2)
    for priority, exp in enumerate(experiments, 1):
        out += strandOverlay(db, "proCapNet", exp, "pred", "/pred", "proCapNet",
                             f"{exp.cell} Predicted",
                             f"{exp.cell} ({exp.modelAcc}) ProCapNet predicted PRO-cap",
                             priority)
    if db == "hg38":
        for priority, exp in enumerate(experiments, 11):
            out += contribStanza(db, exp, priority)
    return out

def encode4ProCapComposite(db, experiments):
    out = compositeStanza("encode4ProCap", "PRO-cap",
                          "PRO-cap nascent RNA transcription start sites from ENCODE 4",
                          experiments, [("procapPos", "PRO-cap_+"),
                                        ("procapNeg", "PRO-cap_-")], 1)
    for priority, exp in enumerate(experiments, 1):
        out += strandOverlay(db, "encode4ProCap", exp, "procap", "", exp.procapAcc,
                             f"{exp.cell} PRO-cap",
                             f"{exp.cell} ({exp.procapAcc}) PRO-cap transcription start"
                             f" sites",
                             priority)
    return out

def releasedSuper(track, shortLabel, longLabel):
    """The released shape of a data source: a top level superTrack holding the
    overlays directly.  The composite shape this file also writes is the one being
    worked on, and it is gated to alpha from trackDb.ra; this is what beta and
    public get until that work lands.  refs #35528"""
    return stanza(0, [
        f"track {track}",
        "superTrack on",
        "group rna",
        f"shortLabel {shortLabel}",
        f"longLabel {longLabel}",
    ])

def releasedOverlay(db, track, exp, dataType, subDir, fileTag, shortLabel, longLabel,
                    priority):
    """One cell line's two strands under the released superTrack.  There is no
    composite to inherit from here, so each overlay carries its own wiggle settings
    and names its own description page."""
    gbdb = GBDB.format(db=db, track=track) + subDir
    out = stanza(0, [
        f"track {track}_{exp.cell}_{dataType}",
        f"superTrack {track} full",
        "container multiWig",
        "aggregate solidOverlay",
        "showSubtrackColorOnUi on",
        "type bigWig",
        "autoScale on",
        "alwaysZero on",
        "maxHeightPixels 100:40:8",
        "windowingFunction maximum",
        "configurable on",
        f"html {track}",
        f"color {exp.color}",
        f"shortLabel {shortLabel}",
        f"longLabel {longLabel}",
        f"priority {priority}",
    ])
    for strand, extra, suffix in (("pos", [], "plus strand"),
                                  ("neg", [f"altColor {exp.color}"], "minus strand")):
        out += stanza(4, [
            f"track {track}_{exp.cell}_{dataType}_{strand}",
            f"parent {track}_{exp.cell}_{dataType}",
            "type bigWig",
            f"bigDataUrl {gbdb}/{exp.cell}.{fileTag}.{strand}.bw",
            f"color {exp.color}",
        ] + extra + [
            f"shortLabel {shortLabel} {'+' if strand == 'pos' else '-'}",
            f"longLabel {longLabel}, {suffix}",
        ])
    return out

def releasedContrib(db, exp, priority):
    gbdb = GBDB.format(db=db, track="proCapNet")
    return stanza(0, [
        f"track proCapNet_{exp.cell}_contrib",
        "superTrack proCapNet hide",
        "type bigWig",
        f"bigDataUrl {gbdb}/contrib/{exp.cell}.proCapNet-contrib.bw",
        "logo on",
        "autoScale on",
        "alwaysZero on",
        "maxHeightPixels 100:40:8",
        "configurable on",
        "html proCapNet",
        f"color {exp.color}",
        f"shortLabel {exp.cell} Contribution",
        f"longLabel {exp.cell} ({exp.modelAcc}) ProCapNet sequence-contribution scores",
        f"priority {priority}",
    ])

def releasedProCapNet(db, experiments):
    out = releasedSuper("proCapNet", "ProCapNet", "ProCapNet predicted PRO-cap")
    for priority, exp in enumerate(experiments, 1):
        out += releasedOverlay(db, "proCapNet", exp, "pred", "/pred", "proCapNet",
                               f"{exp.cell} Predicted",
                               f"{exp.cell} ({exp.modelAcc}) ProCapNet predicted PRO-cap",
                               priority)
    if db == "hg38":
        for priority, exp in enumerate(experiments, 11):
            out += releasedContrib(db, exp, priority)
    return out

def releasedEncode4ProCap(db, experiments):
    out = releasedSuper("encode4ProCap", "PRO-cap",
                        "PRO-cap nascent RNA transcription start sites from ENCODE 4")
    for priority, exp in enumerate(experiments, 1):
        out += releasedOverlay(db, "encode4ProCap", exp, "procap", "", exp.procapAcc,
                               f"{exp.cell} PRO-cap",
                               f"{exp.cell} ({exp.procapAcc}) PRO-cap transcription"
                               f" start sites",
                               priority)
    return out

def proCapNetTrackDb(opts, args):
    """Two files per assembly, one per shape, chosen by the release tag on the
    include line in trackDb.ra.  transcriptionStart.ra is the container and the
    composites inside it, which is the shape being worked on and is alpha until
    #38441 is fixed.  transcriptionStartPublic.ra is the shape already released,
    each data source a top level superTrack holding its overlays."""
    experiments = list(TsvReader(args.experimentsTsv))
    alpha = superStanza()
    released = ""
    if args.db == "hg38":
        alpha += encode4ProCapComposite(args.db, experiments)
        released += releasedEncode4ProCap(args.db, experiments)
    alpha += proCapNetComposite(args.db, experiments)
    released += releasedProCapNet(args.db, experiments)
    writeRa(args.db, args.outDir, "transcriptionStart", alpha)
    writeRa(args.db, args.outDir, "transcriptionStartPublic", released)

def main():
    opts, args = parseArgs()
    with cli.ErrorHandler(noStackExcepts=(OSError, cli.PycbioException)):
        proCapNetTrackDb(opts, args)

main()
