#!/usr/bin/env python3

#######################################################################
#######################################################################
## Created on June 29th 2018 to annotate merged peaks
#######################################################################
#######################################################################

import os
import errno
import argparse

############################################
############################################
## PARSE ARGUMENTS
############################################
############################################

Description = "Add sample boolean files and aggregate columns from merged MACS narrow or broad peak file."
Epilog = """Example usage: python macs2_merged_expand.py <MERGED_INTERVAL_FILE> <SAMPLE_NAME_LIST> <OUTFILE> --is_narrow_peak --min_replicates 1"""

argParser = argparse.ArgumentParser(description=Description, epilog=Epilog)

## REQUIRED PARAMETERS
argParser.add_argument("MERGED_INTERVAL_FILE", help="Merged MACS2 interval file created using linux sort and mergeBed.")
argParser.add_argument(
    "SAMPLE_NAME_LIST",
    help="Comma-separated list of sample names as named in individual MACS2 broadPeak/narrowPeak output file e.g. SAMPLE_R1 for SAMPLE_R1_peak_1.",
)
argParser.add_argument("OUTFILE", help="Full path to output directory.")

## OPTIONAL PARAMETERS
argParser.add_argument(
    "-in",
    "--is_narrow_peak",
    dest="IS_NARROW_PEAK",
    help="Whether merged interval file was generated from narrow or broad peak files (default: False).",
    action="store_true",
)
argParser.add_argument(
    "-mr",
    "--min_replicates",
    type=int,
    dest="MIN_REPLICATES",
    default=1,
    help="Minumum number of replicates per sample required to contribute to merged peak (default: 1).",
)
args = argParser.parse_args()

############################################
############################################
## HELPER FUNCTIONS
############################################
############################################


def makedir(path):
    if not len(path) == 0:
        try:
            os.makedirs(path)
        except OSError as exception:
            if exception.errno != errno.EEXIST:
                raise


############################################
############################################
## MAIN FUNCTION
############################################
############################################

## MergedIntervalTxtFile is file created using commands below:
## 1) broadPeak
## sort -k1,1 -k2,2n <MACS_BROADPEAK_FILES_LIST> | mergeBed -c 2,3,4,5,6,7,8,9 -o collapse,collapse,collapse,collapse,collapse,collapse,collapse,collapse > merged_peaks.txt
## 2) narrowPeak
## sort -k1,1 -k2,2n <MACS_NARROWPEAK_FILE_LIST> | mergeBed -c 2,3,4,5,6,7,8,9,10 -o collapse,collapse,collapse,collapse,collapse,collapse,collapse,collapse,collapse > merged_peaks.txt


def macs2_merged_expand(MergedIntervalTxtFile, SampleNameList, OutFile, isNarrow=False, minReplicates=1):
    makedir(os.path.dirname(OutFile))

    combFreqDict = {}
    totalOutIntervals = 0
    SampleNameList = sorted(SampleNameList)
    fin = open(MergedIntervalTxtFile, "r")
    fout = open(OutFile, "w")
    oFields = (
        ["chr", "start", "end", "interval_id", "num_peaks", "num_samples"]
        + [x + ".bool" for x in SampleNameList]
        + [x + ".fc" for x in SampleNameList]
        + [x + ".qval" for x in SampleNameList]
        + [x + ".pval" for x in SampleNameList]
        + [x + ".start" for x in SampleNameList]
        + [x + ".end" for x in SampleNameList]
    )
    if isNarrow:
        oFields += [x + ".summit" for x in SampleNameList]
    fout.write("\t".join(oFields) + "\n")
    while True:
        line = fin.readline()
        if line:
            lspl = line.strip().split("\t")

            chromID = lspl[0]
            mstart = int(lspl[1])
            mend = int(lspl[2])
            starts = [int(x) for x in lspl[3].split(",")]
            ends = [int(x) for x in lspl[4].split(",")]
            names = lspl[5].split(",")
            fcs = [float(x) for x in lspl[8].split(",")]
            pvals = [float(x) for x in lspl[9].split(",")]
            qvals = [float(x) for x in lspl[10].split(",")]
            summits = []
            if isNarrow:
                summits = [int(x) for x in lspl[11].split(",")]

            ## GROUP SAMPLES BY REMOVING TRAILING *_R*
            groupDict = {}
            for sID in ["_".join(x.split("_")[:-2]) for x in names]:
                gID = "_".join(sID.split("_")[:-1])
                if gID not in groupDict:
                    groupDict[gID] = []
                if sID not in groupDict[gID]:
                    groupDict[gID].append(sID)

            ## GET SAMPLES THAT PASS REPLICATE THRESHOLD
            passRepThreshList = []
            for gID, sIDs in groupDict.items():
                if len(sIDs) >= minReplicates:
                    passRepThreshList += sIDs

            ## GET VALUES FROM INDIVIDUAL PEAK SETS
            fcDict = {}
            qvalDict = {}
            pvalDict = {}
            startDict = {}
            endDict = {}
            summitDict = {}
            for idx in range(len(names)):
                sample = "_".join(names[idx].split("_")[:-2])
                if sample in passRepThreshList:
                    if sample not in fcDict:
                        fcDict[sample] = []
                    fcDict[sample].append(str(fcs[idx]))
                    if sample not in qvalDict:
                        qvalDict[sample] = []
                    qvalDict[sample].append(str(qvals[idx]))
                    if sample not in pvalDict:
                        pvalDict[sample] = []
                    pvalDict[sample].append(str(pvals[idx]))
                    if sample not in startDict:
                        startDict[sample] = []
                    startDict[sample].append(str(starts[idx]))
                    if sample not in endDict:
                        endDict[sample] = []
                    endDict[sample].append(str(ends[idx]))
                    if isNarrow:
                        if sample not in summitDict:
                            summitDict[sample] = []
                        summitDict[sample].append(str(summits[idx]))

            samples = sorted(fcDict.keys())
            if samples != []:
                numSamples = len(samples)
                boolList = ["TRUE" if x in samples else "FALSE" for x in SampleNameList]
                fcList = [";".join(fcDict[x]) if x in samples else "NA" for x in SampleNameList]
                qvalList = [";".join(qvalDict[x]) if x in samples else "NA" for x in SampleNameList]
                pvalList = [";".join(pvalDict[x]) if x in samples else "NA" for x in SampleNameList]
                startList = [";".join(startDict[x]) if x in samples else "NA" for x in SampleNameList]
                endList = [";".join(endDict[x]) if x in samples else "NA" for x in SampleNameList]
                oList = [
                    str(x)
                    for x in [chromID, mstart, mend, "Interval_" + str(totalOutIntervals + 1), len(names), numSamples]
                    + boolList
                    + fcList
                    + qvalList
                    + pvalList
                    + startList
                    + endList
                ]
                if isNarrow:
                    oList += [";".join(summitDict[x]) if x in samples else "NA" for x in SampleNameList]
                fout.write("\t".join(oList) + "\n")

                tsamples = tuple(sorted(samples))
                if tsamples not in combFreqDict:
                    combFreqDict[tsamples] = 0
                combFreqDict[tsamples] += 1
                totalOutIntervals += 1

        else:
            fin.close()
            fout.close()
            break

    ## WRITE FILE FOR INTERVAL INTERSECT ACROSS SAMPLES.
    ## COMPATIBLE WITH UPSETR PACKAGE.
    fout = open(OutFile[:-4] + ".intersect.txt", "w")
    combFreqItems = sorted([(combFreqDict[x], x) for x in combFreqDict.keys()], reverse=True)
    for k, v in combFreqItems:
        fout.write("%s\t%s\n" % ("&".join(v), k))
    fout.close()


############################################
############################################
## RUN FUNCTION
############################################
############################################

macs2_merged_expand(
    MergedIntervalTxtFile=args.MERGED_INTERVAL_FILE,
    SampleNameList=args.SAMPLE_NAME_LIST.split(","),
    OutFile=args.OUTFILE,
    isNarrow=args.IS_NARROW_PEAK,
    minReplicates=args.MIN_REPLICATES,
)

############################################
############################################
############################################
############################################
