import os
from os import path as osp
import glob
import shutil
import tempfile
from pycbio.sys import fileOps
import pipettor

asmName = "cenSat1"
os.environ["PATH"] = "../../bin:" + os.environ["PATH"]

##
# IMPORTANT
# must create hub/cenSat1/cenSat1.seqs with the list of contigs
##

resultsDir = "/nanopore/markd/t2t-centromere/T2T_cenSat"

##
# output directories
##
hubDir = osp.join(resultsDir, "hub")
asmHubDir = osp.join(hubDir, asmName)
asmDataDir = osp.join(asmHubDir, "data")

cenSatWgDir = "/public/groups/nanopore/karen/chm13/cenSat_WG"
asmSrcDir = osp.join(cenSatWgDir, "assembly")
phase1SrcDir = osp.join(cenSatWgDir, "phase1/browser")

##
# hub files
##
hubSrcDir = "../../hub"
def findHubFiles():
    prelen = len(Path(hubSrcDir).parts)
    return [osp.join(*p.parts[prelen:])
            for  p in (list(Path(hubSrcDir).rglob("*.txt"))
                       + list(Path(hubSrcDir).rglob("*.html")))]
hubFiles = findHubFiles()
hubInFiles = [osp.join(hubSrcDir, f) for f in hubFiles]
hubOutFiles = [osp.join(hubDir, f) for f in hubFiles]

##
# inputs (need in multiple places)
##
asmSeqs = osp.join(asmHubDir, asmName + ".seqs")
contigs = list(fileOps.readFileLines(asmSeqs))

##
# results
##
asmTwoBit = osp.join(asmHubDir, asmName + ".2bit")
asmSizes = osp.join(asmHubDir, asmName + ".sizes")
centroRegionsBigBed = osp.join(asmHubDir, "centroRegions.bigBed")
rmskBigBed = osp.join(asmHubDir, "rmsk.bigBed")
horArrayBigBed = osp.join(asmHubDir, "horArray.bigBed")
horHmmBigBed = osp.join(asmHubDir, "horHmm.bigBed")
allGenomeSFBigBed = osp.join(asmHubDir, "allGenomeSF.bigBed")
sf1HorBigBed = osp.join(asmHubDir, "sf1Hor.bigBed")
pubMonBigBed = osp.join(asmHubDir, "pubMon.bigBed")
hg38CenSat1BigBed = osp.join(asmHubDir, "hg38CenSat1.bigBed")
cenSat1Hg38BigBed = osp.join(asmHubDir, "cenSat1Hg38.bigBed")
monDecompositionBigBed = osp.join(asmHubDir, "monDecomposition.bigBed")
rawMethCallsBigWig = osp.join(asmHubDir, "rawMethCalls.bigWig")
hsat10kCapBigBed = osp.join(asmHubDir, "hsat10kCap.bigBed")
hsatThreshold10kCapBigBed = osp.join(asmHubDir, "hsatThreshold10kCap.bigBed")
hsatThreshold10kCapOriBigBed = osp.join(asmHubDir, "hsatThreshold10kCapOri.bigBed")
asHorAnnBigBed = osp.join(asmHubDir, "asHorAnn.bigBed")
asHorAnnOriBigBed = osp.join(asmHubDir, "asHorOriAnn.bigBed")
horDecompositionBigBed = osp.join(asmHubDir, "horDecomposition.bigBed")
hicanuSegDupsBigBed = osp.join(asmDataDir, "hicanuSegDupsBigBed.bigBed")

allResults = [asmTwoBit, asmSizes, centroRegionsBigBed,
              rmskBigBed, horArrayBigBed, horHmmBigBed, allGenomeSFBigBed,
              sf1HorBigBed, pubMonBigBed, hg38CenSat1BigBed, cenSat1Hg38BigBed,
              monDecompositionBigBed, hsat10kCapBigBed, hsatThreshold10kCapBigBed,
              hsatThreshold10kCapOriBigBed,
              asHorAnnBigBed, asHorAnnOriBigBed,
              horDecompositionBigBed, hicanuSegDupsBigBed] + hubOutFiles
# BROKEN: rawMethCallsBigWig

##
# misc defs
##
bedSortCmd = ["sort", "-k1,1", "-k2,2n"]
showCoords = "../../../packages/MUMmer3.23/show-coords"

##
# debugging
##
if True:
    from pycbio.sys import loggingOps
    import logging
    loggingOps.setupStderrLogger("DEBUG")
    pipettor.setDefaultLogger(logging.getLogger())
bigKeepTmp = False

##
# functions to do common stuff
##
def inputToLst(inSpec):
    if isinstance(inSpec, str):
        return [inSpec]
    elif isinstance(inSpec, set):
        if isinstance(next(iter(inSpec)), list):
            # set of lists
            specs = []
            for l in inSpec:
                specs += l
            return specs
        else:
            return list(inSpec)
    elif isinstance(inSpec, (list, tuple)):
        return inSpec
    else:
        raise Exception("Don't know how to make input into a list: {}: {}".format(type(inSpec), inSpec))


def buildBigBed(inBeds, outBigBed, *, bedType="bed4", asFile=None, filterCmd=None):
    inBeds = inputToLst(inBeds)
    with tempfile.NamedTemporaryFile(delete=not bigKeepTmp) as bedFh:
        # discard non-cenSat contigs
        if str(inBeds[0]).endswith(".gz"):
            cmd = [["zcat"] + inBeds]
        else:
            cmd = [["cat"] + inBeds]
        cmd += [["selectById", 1, asmSeqs, 1], bedSortCmd]
        if filterCmd is not None:
            cmd.append(filterCmd)
        pipettor.run(cmd, stdout=bedFh.name)
        bbcmd = ["bedToBigBed", "-type=" + bedType, "-tab"]
        if asFile is not None:
            bbcmd += ["-as=" + asFile]
        bbcmd += [bedFh.name, asmSizes, outBigBed]
        pipettor.run(bbcmd)

def buildBigWig(inWigs, outBigWig, *, clip=False):
    if isinstance(inWigs, str):
        inWigs = [inWigs]
    with tempfile.NamedTemporaryFile(delete=not bigKeepTmp) as wigFh:
        pipettor.run(["cat"] + inWigs, stdout=wigFh.name)
        opts = []
        if clip:
            opts.append("-clip")
        pipettor.run(["wigToBigWig"] + opts + [wigFh.name, asmSizes, outBigWig])

###
# rule that drives everything
###

rule all:
    input: allResults
    shell: "hubCheck https://courtyard.gi.ucsc.edu/~markd/T2T_cenSat/hub/hub.txt"

rule clean:
    run:
        fileOps.rmFiles(*allResults)



###
# genome file
###
genomeFasta = osp.join(asmSrcDir, "chm13_hicanu_hifi_20k.fasta")

rule genomeSeqs:
    input: genomeFasta, asmSeqs
    output: asmTwoBit, asmSizes
    run:
        pipettor.run([("faSomeRecords", genomeFasta, asmSeqs, "/dev/stdout"),
                      ("faToTwoBit", "/dev/stdin", asmTwoBit)])
        pipettor.run(["twoBitInfo", asmTwoBit, asmSizes])


###
# copy hub files
###
rule hubText:
    input: hubInFiles
    output: hubOutFiles
    run:
        for i in range(len(input)):
            shutil.copyfile(input[i], output[i])

###
# centromere region track
###
centroRegionBed = osp.join(phase1SrcDir, "T2T_cenSat_Phase1_2_Array_Annotation.cenRegions.tsv.bed")
rule centroRegionsTrack:
    input: centroRegionBed, asmSizes
    output: centroRegionsBigBed
    run:
        buildBigBed(centroRegionBed, centroRegionsBigBed, bedType="bed9")


##
# repeat masker
##
rmskDir = osp.join(asmSrcDir, "rm")
rmskOutFiles = []
for tig in contigs:
    pat = osp.join(rmskDir, tig + ".out/*.fasta.out")
    touts = glob.glob(pat)
    if len(touts) == 0:
        pass
        #print("WARNING: no repeats for", tig, file=sys.stderr)
    rmskOutFiles.extend(touts)

rule rmsk:
    input: rmskOutFiles, asmSeqs
    output: rmskBigBed
    run:
        # regions are named in the form tig00018121:0-1387635,
        # but cover the whole contig
        with tempfile.NamedTemporaryFile() as bedFh:
            for f in rmskOutFiles:
                pipettor.run([["RM2Bed.py", "--stdout", "--log_level=WARN", f],
                              ["editRepeatMaskerBed"]], stdout=bedFh)
            buildBigBed(bedFh.name, rmskBigBed, bedType="bed9")

##
# HOR arrays, etc
##
horArrayBed = osp.join(phase1SrcDir, "T2T_cenSat_Phase1_2_Array_Annotation.cenSat.tsv.2.bed")
rule horArray:
    input: horArrayBed, asmSizes
    output: horArrayBigBed
    run:
        buildBigBed(horArrayBed, horArrayBigBed, bedType="bed9")

horHmmBed = osp.join(phase1SrcDir, "cenSat_HOR_hmmOut.bed")
rule horHmm:
    input: horHmmBed, asmSizes
    output: horHmmBigBed
    run:
        buildBigBed(horHmmBed, horHmmBigBed, filterCmd=["cut", "-f", "1-4"])

##
# SF tracks
##

allGenomeSFBed = osp.join(phase1SrcDir, "all_genome_SF_track.chm13_hicanu_hifi_20k.bed")
rule allGenomeSF:
    input: allGenomeSFBed, asmSizes
    output: allGenomeSFBigBed
    run:
        buildBigBed(allGenomeSFBed, allGenomeSFBigBed, bedType="bed9")

sf1HorBed = osp.join(phase1SrcDir, "SF1_HOR_track.chm13_hicanu_hifi_20k.bed")
rule sf1Hor:
    input: sf1HorBed, asmSizes
    output: sf1HorBigBed
    run:
        buildBigBed(sf1HorBed, sf1HorBigBed, bedType="bed9")

pubMonBed = osp.join(phase1SrcDir, "pubMon_track.chm13_hicanu_hifi_20k.bed")
rule pubMon:
    input: pubMonBed, asmSizes
    output: pubMonBigBed
    run:
        buildBigBed(pubMonBed, pubMonBigBed, bedType="bed9")

##
# nunmer alignments
##
tigHg38Delta = osp.join(asmDataDir, "chm13_20k_hiFi_master0103.dnadiff/asm.contigs.fasta.againstPrimary.delta.diff.1delta")
hg38ChromAlias = osp.join(asmDataDir, "hg38ChromAlias.tsv")
rule hg38CenSat1Align:
    input: tigHg38Delta, hg38ChromAlias
    output: hg38CenSat1BigBed
    run:
        with tempfile.NamedTemporaryFile() as bedFh:
            pipettor.run([[showCoords, "-T", tigHg38Delta],
                          ["nucmerToBed", "--chromAliasTsv=" + hg38ChromAlias, "seq1"]], stdout=bedFh.name)
            buildBigBed(bedFh.name, hg38CenSat1BigBed, bedType="bed12")

rule cenSat1Hg38Align:
    input: tigHg38Delta, hg38ChromAlias
    output: cenSat1Hg38BigBed
    run:
        with tempfile.NamedTemporaryFile() as bedFh:
            pipettor.run([[showCoords, "-T", tigHg38Delta],
                          ["nucmerToBed", "--chromAliasTsv=" + hg38ChromAlias, "seq2"]], stdout=bedFh.name)
            buildBigBed(bedFh.name, cenSat1Hg38BigBed, bedType="bed12")

##
# monDecomposition
monDecompBeds = list(glob.glob(osp.join(phase1SrcDir, "042720_MonDecomposition/decomposition*.bed")))
rule monDecomp:
    input: monDecompBeds
    output: monDecompositionBigBed
    run:
        buildBigBed(monDecompBeds, monDecompositionBigBed, bedType="bed9",
                    filterCmd=["fixBedScore"])

#
# raw meth calls
rawMethCallsWigs = glob.glob(osp.join(phase1SrcDir, "tig*_raw_meth_calls.wig"))
rule rawMethCalls:
    input: rawMethCallsWigs
    output: rawMethCallsBigWig
    run:
        # badly broken
        # had error: chromosome tig00001118 has 126714748 bases, but item ends at 126714756, so clip
        buildBigWig(rawMethCallsWigs, rawMethCallsBigWig, clip=True)

##
# HSAT annotation
rule hsat10kCap:
    input: osp.join(asmDataDir, "hsat-annotation/HSat23_BrowserTrack_v1_10kCap.bed.gz")
    output: hsat10kCapBigBed
    run:
        buildBigBed({input}, hsat10kCapBigBed, bedType="bed9",
                    filterCmd=["fixBedScore"])

rule hsatThreshold10kCap:
    input: osp.join(asmDataDir, "hsat-annotation/HSat23_BrowserTrack_v1_threshold_10kCap.bed.gz")
    output: hsatThreshold10kCapBigBed
    run:
        buildBigBed({input}, hsatThreshold10kCapBigBed, bedType="bed9",
                    filterCmd=["fixBedScore"])

rule hsatThreshold10kCapOri:
    input: osp.join(phase1SrcDir, "HSat23_BrowserTrack_v1_threshold_10kCap.oriColor.bed")
    output: hsatThreshold10kCapOriBigBed
    run:
        buildBigBed({input}, hsatThreshold10kCapOriBigBed, bedType="bed9",
                    filterCmd=["fixBedScore"])

# AS HOR annotation
rule asHorAnn:
    input: osp.join(phase1SrcDir, "chm13_hicanu_hifi_20k-vs-AS-HORs-hmmer3.0-22042018.hmm-tbl.bed.gz")
    output: asHorAnnBigBed
    run:
        buildBigBed({input}, asHorAnnBigBed, bedType="bed9",
                    filterCmd=["fixBedScore"])

rule asHorAnnOri:
    input: osp.join(phase1SrcDir, "chm13_hicanu_hifi_20k-vs-AS-HORs-hmmer3.0-22042018.hmm-tbl.ori.bed")
    output: asHorAnnOriBigBed
    run:
        buildBigBed({input}, asHorAnnOriBigBed, bedType="bed9",
                    filterCmd=["fixBedScore"])

# horDecomposition
rule horDecomposition:
    input: glob.glob(osp.join(phase1SrcDir, "042120_HORDecomposition/HORDecompositionNew/*bed"))
    output: horDecompositionBigBed
    run:
        buildBigBed({input}, horDecompositionBigBed, bedType="bed9")

rule hicanuSegDups:
    input:  osp.join(resultsDir, "cenSat1/chm13_hicanu_hifi_20k_no_bubbles_sedef_out/final.browser.bed.gz")
    output: hicanuSegDupsBigBed
    run:
        buildBigBed({input}, hicanuSegDupsBigBed, bedType="bed9+", asFile="hicanuSegDups.as")
